{"cells":[{"cell_type":"markdown","metadata":{"id":"-m4sB0AwIO_X"},"source":["# Recitation 0: Weights and Biases\n","\n","In this recitation, you will learn about the importance of performance visualization and model tracking using [WandB](https://wandb.ai/), a tool for performance visualization, model and data version controlling and hyperparameter tuning."]},{"cell_type":"markdown","metadata":{"id":"nBa-j_M5sBHX"},"source":["## Installation and Libraries"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"8nqdSi0QvUHQ"},"outputs":[],"source":["## Installing WandB\n","!pip install wandb -qqq"]},{"cell_type":"code","execution_count":null,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"WCGHNaJZqzix","outputId":"8b247b77-3761-4ebb-d7e6-58fe6aac4951","executionInfo":{"status":"ok","timestamp":1734648029334,"user_tz":-120,"elapsed":22115,"user":{"displayName":"Floris Nzabakira","userId":"16890065093683199713"}}},"outputs":[{"output_type":"stream","name":"stdout","text":["Device: cuda\n"]}],"source":["import torch\n","import torch.nn as nn\n","import torch.optim as optim\n","from torchvision import datasets\n","from torchvision.transforms import ToTensor\n","\n","device = 'cuda' if torch.cuda.is_available() else 'cpu'\n","print(\"Device: \", device)\n","\n","from tqdm import tqdm"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"7BoEvflSsHMK","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1734648035766,"user_tz":-120,"elapsed":6439,"user":{"displayName":"Floris Nzabakira","userId":"16890065093683199713"}},"outputId":"6c076261-d260-4690-e40d-358e00e2eef9"},"outputs":[{"output_type":"stream","name":"stderr","text":["\u001b[34m\u001b[1mwandb\u001b[0m: Using wandb-core as the SDK backend. Please refer to https://wandb.me/wandb-core for more information.\n","\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mnzafloris\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"]},{"output_type":"execute_result","data":{"text/plain":["True"]},"metadata":{},"execution_count":3}],"source":["import wandb, os\n","os.environ['WANDB_API_KEY'] = \"\" #your key here\n","wandb.login()"]},{"cell_type":"markdown","metadata":{"id":"STaDx_PhsY5S"},"source":["## Helper functions and Model"]},{"cell_type":"code","execution_count":null,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"fRjtOHvzq8D0","outputId":"e74d762e-69c2-4962-8e67-2ee704d1e621","executionInfo":{"status":"ok","timestamp":1734648044879,"user_tz":-120,"elapsed":9128,"user":{"displayName":"Floris Nzabakira","userId":"16890065093683199713"}}},"outputs":[{"output_type":"stream","name":"stdout","text":["Downloading https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz to data/cifar-10-python.tar.gz\n"]},{"output_type":"stream","name":"stderr","text":["100%|██████████| 170M/170M [00:03<00:00, 46.9MB/s]\n"]},{"output_type":"stream","name":"stdout","text":["Extracting data/cifar-10-python.tar.gz to data\n","Files already downloaded and verified\n"]}],"source":["data_train = datasets.CIFAR10(\n"," root = 'data',\n"," train = True,\n"," transform = ToTensor(),\n"," download = True,\n",")\n","data_test = datasets.CIFAR10(\n"," root = 'data',\n"," train = False,\n"," download = True,\n"," transform = ToTensor()\n",")"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"thxzbhQ-rAzL"},"outputs":[],"source":["def build_data(batch_size, data_train, data_test):\n"," train_loader = torch.utils.data.DataLoader(data_train, batch_size=batch_size, shuffle=True)\n"," test_loader = torch.utils.data.DataLoader(data_test, batch_size=batch_size, shuffle=False)\n"," return train_loader, test_loader"]},{"cell_type":"code","execution_count":null,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"JD8Q9XcbtALD","outputId":"78c85fe8-210e-407b-e43d-c19efe7e0f94","executionInfo":{"status":"ok","timestamp":1734648044880,"user_tz":-120,"elapsed":13,"user":{"displayName":"Floris Nzabakira","userId":"16890065093683199713"}}},"outputs":[{"output_type":"stream","name":"stdout","text":["Network(\n"," (CNN): Sequential(\n"," (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(2, 2))\n"," (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n"," (2): ReLU()\n"," (3): AvgPool2d(kernel_size=9, stride=9, padding=0)\n"," (4): Flatten(start_dim=1, end_dim=-1)\n"," )\n"," (classification): Linear(in_features=576, out_features=10, bias=True)\n",")\n"]}],"source":["class Network(nn.Module):\n","\n"," def __init__(self):\n","\n"," super(Network, self).__init__()\n","\n"," self.CNN = nn.Sequential(\n"," nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=2),\n"," nn.BatchNorm2d(64),\n"," nn.ReLU(),\n","\n"," nn.AvgPool2d(kernel_size=9),\n"," nn.Flatten()\n"," )\n","\n"," self.classification = nn.Linear(576, 10)\n"," def forward(self, x):\n","\n"," x_cnn = self.CNN(x)\n"," res = self.classification(x_cnn)\n","\n"," return res\n","\n","model = Network().to(device)\n","print(model)"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"5Pg2TkSyEXXp","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1734648045308,"user_tz":-120,"elapsed":437,"user":{"displayName":"Floris Nzabakira","userId":"16890065093683199713"}},"outputId":"7a4d282a-6789-4078-c977-cba07c032ad5"},"outputs":[{"output_type":"execute_result","data":{"text/plain":["torch.Size([64, 10])"]},"metadata":{},"execution_count":7}],"source":["train_loader, test_loader = build_data(64, data_train, data_test)\n","\n","for x, y in train_loader:\n"," break\n","model(x.to(device)).shape"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"APS6OULkPbMJ"},"outputs":[],"source":["def get_optim(optimizer, learning_rate, model):\n"," if optimizer=='sgd':\n"," return optim.SGD(model.parameters(), lr=learning_rate)\n"," else:\n"," return optim.Adam(model.parameters(), lr=learning_rate)"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"WWLpSANVs1EH"},"outputs":[],"source":["def train_epoch(model, loader, optimizer, criterion, scaler):\n"," num_correct = 0\n"," total_loss = 0\n","\n"," for i, (x, y) in enumerate(loader):\n"," optimizer.zero_grad()\n","\n"," x = x.cuda()\n"," y = y.cuda()\n","\n"," with torch.cuda.amp.autocast():\n"," outputs = model(x)\n"," loss = criterion(outputs, y)\n","\n"," total_loss += float(loss)\n","\n"," scaler.scale(loss).backward()\n"," scaler.step(optimizer)\n"," scaler.update()\n"," ep_loss = float(total_loss / len(loader))\n","\n"," return model, ep_loss"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"BMJlPL5CYPFM"},"outputs":[],"source":["def train(model, finish= True):\n","\n"," # Dont worry about all this, you'll be very familiar with it after HW1\n","\n"," best_acc = 0\n","\n"," for epoch in range(run_config['epochs']):\n"," batch_bar = tqdm(total=len(train_loader), dynamic_ncols=True, leave=False, position=0, desc='Train')\n","\n"," num_correct = 0\n"," total_loss = 0\n","\n"," for i, (x, y) in enumerate(train_loader):\n"," optimizer.zero_grad()\n","\n"," x = x.cuda()\n"," y = y.cuda()\n","\n"," with torch.cuda.amp.autocast():\n"," outputs = model(x)\n"," loss = criterion(outputs, y)\n","\n"," num_correct += int((torch.argmax(outputs, axis=1) == y).sum())\n"," total_loss += float(loss)\n","\n"," batch_bar.set_postfix(\n"," acc=\"{:.04f}%\".format(100 * num_correct / ((i + 1) * run_config['batch_size'])),\n"," loss=\"{:.04f}\".format(float(total_loss / (i + 1))),\n"," num_correct=num_correct,\n"," lr=\"{:.04f}\".format(float(optimizer.param_groups[0]['lr'])))\n","\n"," scaler.scale(loss).backward()\n"," scaler.step(optimizer)\n"," scaler.update()\n","\n","\n"," batch_bar.update()\n"," batch_bar.close()\n","\n"," train_loss = float(total_loss / len(train_loader))\n"," train_acc = 100 * num_correct / (len(train_loader) * run_config['batch_size'])\n"," lr = float(optimizer.param_groups[0]['lr'])\n","\n"," print(\"Epoch {}/{}: Train Acc {:.04f}%, Train Loss {:.04f}, Learning Rate {:.04f}\".format(\n"," epoch + 1,\n"," run_config['epochs'],\n"," train_acc ,\n"," train_loss,\n"," lr\n"," )\n"," )\n","\n"," # What to log\n","\n"," metrics = {\n"," \"train_loss\":train_loss,\n"," \"train_acc\": train_acc,\n"," 'lr': lr\n"," }\n","\n"," # Log to run\n"," wandb.log(metrics)\n","\n"," # Updating the model version\n","\n"," if train_acc > best_acc:\n"," best_acc = train_acc\n","\n"," # Saving the model and optimizer states\n","\n"," torch.save({\n"," 'model_state_dict': model.state_dict(),\n"," 'optimizer_state_dict': optimizer.state_dict()\n"," }, \"Model.pth\")\n","\n"," # ALTERNATIVE 1: Saving Files as Artifacts\n"," # Creating Artifact\n"," model_artifact = wandb.Artifact(run_config['model'], type='model')\n","\n"," # Adding model file to Artifact\n"," model_artifact.add_file(\"Model.pth\")\n","\n"," # Saving Artifact to WandB\n"," run.log_artifact(model_artifact)\n","\n"," # ALTERNATIVE 2: Saving Files as Files\n"," wandb.save(\"Model.pth\")\n","\n"," if finish:\n"," wandb.finish()"]},{"cell_type":"markdown","metadata":{"id":"8IoBEQ9FYTcT"},"source":["## Simple Usage"]},{"cell_type":"markdown","metadata":{"id":"bBdRpZygU8Xk"},"source":["You can run the training function and log the performance metrics of your choice into the WandB GUI. This simple method will allow you to monitor trends in a specefic run configuration as well as comparing different runs"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"pCP4X9rJYoiT","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1734648045310,"user_tz":-120,"elapsed":23,"user":{"displayName":"Floris Nzabakira","userId":"16890065093683199713"}},"outputId":"93839d2c-753b-4fd4-b103-ef7ad46f6089"},"outputs":[{"output_type":"stream","name":"stderr","text":[":15: FutureWarning: `torch.cuda.amp.GradScaler(args...)` is deprecated. Please use `torch.amp.GradScaler('cuda', args...)` instead.\n"," scaler = torch.cuda.amp.GradScaler()\n"]}],"source":["run_config = {\n"," 'model': '1-2dcnn',\n"," 'optimizer':'sgd',\n"," 'lr': 2e-3,\n"," 'batch_size':64,\n"," 'epochs': 5\n","}\n","\n","train_loader, test_loader = build_data(run_config['batch_size'], data_train, data_test)\n","\n","optimizer = get_optim(run_config['optimizer'], run_config['lr'], model)\n","\n","criterion = nn.CrossEntropyLoss()\n","\n","scaler = torch.cuda.amp.GradScaler()"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"90SpRavzYsWD","colab":{"base_uri":"https://localhost:8080/","height":121},"executionInfo":{"status":"ok","timestamp":1734648692091,"user_tz":-120,"elapsed":3318,"user":{"displayName":"Floris Nzabakira","userId":"16890065093683199713"}},"outputId":"92585a24-f8a0-41a5-8978-80c773da5854"},"outputs":[{"output_type":"stream","name":"stderr","text":["\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mnzafloris\u001b[0m (\u001b[33mwandb-starter\u001b[0m). Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"]},{"output_type":"display_data","data":{"text/plain":[""],"text/html":["Tracking run with wandb version 0.19.1"]},"metadata":{}},{"output_type":"display_data","data":{"text/plain":[""],"text/html":["Run data is saved locally in /content/wandb/run-20241219_225136-k7cu7hle"]},"metadata":{}},{"output_type":"display_data","data":{"text/plain":[""],"text/html":["Syncing run 1-2dcnn to Weights & Biases (docs)
"]},"metadata":{}},{"output_type":"display_data","data":{"text/plain":[""],"text/html":[" View project at https://wandb.ai/wandb-starter/wandb-quickstart"]},"metadata":{}},{"output_type":"display_data","data":{"text/plain":[""],"text/html":[" View run at https://wandb.ai/wandb-starter/wandb-quickstart/runs/k7cu7hle"]},"metadata":{}}],"source":["run = wandb.init(\n"," entity=\"wandb-starter\",\n"," project=\"wandb-quickstart\",\n"," job_type=\"model-training\",\n"," name=run_config['model'],\n"," config=run_config\n"," )"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"Yze2NvQAYwQc","colab":{"base_uri":"https://localhost:8080/","height":518},"executionInfo":{"status":"ok","timestamp":1734648768597,"user_tz":-120,"elapsed":73326,"user":{"displayName":"Floris Nzabakira","userId":"16890065093683199713"}},"outputId":"4048d5e5-0e16-4bc8-8783-1c448827ff78"},"outputs":[{"output_type":"stream","name":"stderr","text":["Train: 0%| | 0/782 [00:00:19: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n"," with torch.cuda.amp.autocast():\n"]},{"output_type":"stream","name":"stdout","text":["Epoch 1/5: Train Acc 38.2393%, Train Loss 1.7846, Learning Rate 0.0020\n"]},{"output_type":"stream","name":"stderr","text":[]},{"output_type":"stream","name":"stdout","text":["Epoch 2/5: Train Acc 39.0046%, Train Loss 1.7628, Learning Rate 0.0020\n"]},{"output_type":"stream","name":"stderr","text":[]},{"output_type":"stream","name":"stdout","text":["Epoch 3/5: Train Acc 39.7438%, Train Loss 1.7442, Learning Rate 0.0020\n"]},{"output_type":"stream","name":"stderr","text":[]},{"output_type":"stream","name":"stdout","text":["Epoch 4/5: Train Acc 40.4092%, Train Loss 1.7260, Learning Rate 0.0020\n"]},{"output_type":"stream","name":"stderr","text":[]},{"output_type":"stream","name":"stdout","text":["Epoch 5/5: Train Acc 40.8608%, Train Loss 1.7105, Learning Rate 0.0020\n"]},{"output_type":"display_data","data":{"text/plain":[""],"text/html":[]},"metadata":{}},{"output_type":"display_data","data":{"text/plain":[""],"text/html":["

Run history:


lr▁▁▁▁▁
train_acc▁▃▅▇█
train_loss█▆▄▂▁

Run summary:


lr0.002
train_acc40.86077
train_loss1.71053

"]},"metadata":{}},{"output_type":"display_data","data":{"text/plain":[""],"text/html":[" View run 1-2dcnn at: https://wandb.ai/wandb-starter/wandb-quickstart/runs/k7cu7hle
View project at: https://wandb.ai/wandb-starter/wandb-quickstart
Synced 5 W&B file(s), 0 media file(s), 10 artifact file(s) and 1 other file(s)"]},"metadata":{}},{"output_type":"display_data","data":{"text/plain":[""],"text/html":["Find logs at: ./wandb/run-20241219_225136-k7cu7hle/logs"]},"metadata":{}}],"source":["train(model)"]},{"cell_type":"markdown","metadata":{"id":"i9Iswf4vCE5v"},"source":["## Resume a previous run"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"A8HCCvMx-qL6"},"outputs":[],"source":["RESUME_LOGGING = True ### Change to true to test the code for resuming an existing run"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"uPm3IsTvB13U","colab":{"base_uri":"https://localhost:8080/","height":121},"executionInfo":{"status":"ok","timestamp":1734625214987,"user_tz":-120,"elapsed":2977,"user":{"displayName":"Floris Nzabakira","userId":"16890065093683199713"}},"outputId":"6d2497cf-e8eb-4d92-b823-6e582d62c1a0"},"outputs":[{"output_type":"display_data","data":{"text/plain":[""],"text/html":["Tracking run with wandb version 0.19.1"]},"metadata":{}},{"output_type":"display_data","data":{"text/plain":[""],"text/html":["Run data is saved locally in /content/wandb/run-20241219_162020-n8cxgzyc"]},"metadata":{}},{"output_type":"display_data","data":{"text/plain":[""],"text/html":["Resuming run 1-2dcnn to Weights & Biases (docs)
"]},"metadata":{}},{"output_type":"display_data","data":{"text/plain":[""],"text/html":[" View project at https://wandb.ai/nzafloris/wandb-quickstart"]},"metadata":{}},{"output_type":"display_data","data":{"text/plain":[""],"text/html":[" View run at https://wandb.ai/nzafloris/wandb-quickstart/runs/n8cxgzyc"]},"metadata":{}},{"output_type":"stream","name":"stdout","text":["/content/wandb/run-20241219_162020-n8cxgzyc/files\n"]}],"source":["if RESUME_LOGGING:\n"," run_id = \"4jtvb0pu\" ### Replace with run id string\n"," run = wandb.init(\n"," id = run_id, ### Insert specific run id here if you want to resume a previous run\n"," resume = \"must\", ### You need this to resume previous runs, but comment out reinit = True when using this\n"," project = \"wandb-quickstart\", ### Project should be created in your wandb account\n"," )\n","\n"," print(run.dir)\n"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"WT49FQAS-qL7"},"outputs":[],"source":["### Test code to try appending metrics to previously logged metrics in the run\n","### Uncomment to try out\n","\n","# test_new_metrics = {\n","# \"train_loss\":1.5,\n","# \"train_acc\": 40,\n","# 'lr': 0.001\n","# }\n","\n","# wandb.log(test_new_metrics)"]},{"cell_type":"markdown","metadata":{"id":"KrHFe6BKkT9O"},"source":["## HyperParameter Sweeps\n"]},{"cell_type":"markdown","metadata":{"id":"j9pxaI3xVA6h"},"source":["[Sweeps](https://docs.wandb.ai/guides/sweeps) are a way of automating hyperparameter tuning in Deep Learning Models. You set up the values that you want your sweep to try and then check the affect of changing each parameter on each value on the model."]},{"cell_type":"code","execution_count":null,"metadata":{"id":"dfCIa76lkXv0"},"outputs":[],"source":["# Initialize the sweep and set the method (grid, random or bayes\"ian\")\n","\n","sweep_config = {\n"," 'method': 'random'\n"," }"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"5A3bEWfQkbMw"},"outputs":[],"source":["# What is the objective of the sweep (minimize loss, maximize accuracy)\n","\n","metric = {\n"," 'name':'loss',\n"," 'goal':'minimize'\n","}\n","sweep_config['metric'] = metric"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"oXHd5ccNkenT"},"outputs":[],"source":["# Hyperparameters to work with\n","\n","parameters_dict = {\n"," 'optimizer':{\n"," 'values': ['sgd', 'adam']\n"," },\n"," 'learning_rate':{\n"," 'distribution':'uniform',\n"," 'min':2e-4,\n"," 'max':1e-1\n"," },\n"," 'batch_size': {\n"," 'distribution': 'q_log_uniform_values',\n"," 'q':4,\n"," 'min': 16,\n"," 'max': 128\n"," },\n"," 'epochs':{\n"," 'value': 5\n"," }\n","}\n","sweep_config['parameters'] = parameters_dict"]},{"cell_type":"code","execution_count":null,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"vbVvbWECkg3C","outputId":"f83d172f-5d72-4146-8dfe-1b814a9f8e1f","executionInfo":{"status":"ok","timestamp":1734525612440,"user_tz":-120,"elapsed":17,"user":{"displayName":"Floris Nzabakira","userId":"16890065093683199713"}}},"outputs":[{"output_type":"stream","name":"stdout","text":["Create sweep with ID: tz3x5p3x\n","Sweep URL: https://wandb.ai/nzafloris/CIFAR-Sweep2/sweeps/tz3x5p3x\n"]}],"source":["# Initalizing the sweep\n","\n","sweep_id = wandb.sweep(sweep_config, project=\"CIFAR-Sweep2\")"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"ee2J69kxs2pL"},"outputs":[],"source":["def train_sweep(config = None):\n"," with wandb.init(config=config) as run:\n"," run.name=f\"Jeel_{wandb.config.learning_rate}_{wandb.config.batch_size}_{wandb.config.optimizer}\"\n"," config = wandb.config\n","\n"," train_loader, test_loader = build_data(config.batch_size, data_train, data_test)\n","\n"," model = Network().to(device)\n","\n"," optimizer = get_optim(config.optimizer, config.learning_rate, model)\n","\n"," criterion = nn.CrossEntropyLoss()\n","\n"," scaler = torch.cuda.amp.GradScaler()\n","\n"," for epoch in range(config.epochs):\n","\n"," model, loss = train_epoch(model, train_loader, optimizer, criterion, scaler)\n","\n"," wandb.log({'loss': loss})"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"YOHXfLi5kmNg"},"outputs":[],"source":["# Running the sweep\n","\n","wandb.agent(sweep_id, train_sweep, count=2)"]},{"cell_type":"markdown","metadata":{"id":"Q5189Q0WkYh9"},"source":["## Artifact and Model Versioning"]},{"cell_type":"markdown","metadata":{"id":"VwtzDk8gVG5r"},"source":["Artifacts are a method of managing versions for data and models. You can use the artifacts to keep and compare versions of your model while training making it easier to share data and models between team members."]},{"cell_type":"code","execution_count":null,"metadata":{"id":"8jefWW9Qux5c"},"outputs":[],"source":["run_config = {\n"," 'model': '1-2dcnn',\n"," 'optimizer':'adam',\n"," 'lr': 5e-3,\n"," 'batch_size':20,\n"," 'epochs': 5\n","}\n","\n","train_loader, test_loader = build_data(run_config['batch_size'], data_train, data_test)\n","optimizer = get_optim(run_config['optimizer'], run_config['lr'], model)\n","criterion = nn.CrossEntropyLoss()\n","scaler = torch.cuda.amp.GradScaler()"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"hDSSl7WaBQwi"},"outputs":[],"source":["run = wandb.init(\n"," project=\"wandb-quickstart\",\n"," job_type=\"model-training\",\n"," name=run_config['model'],\n"," config=run_config\n"," )"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"t9FjIkHGZdw6"},"outputs":[],"source":["train(model,finish= False) #run should not finish for using artifact"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"29H1tmQjGKsp"},"outputs":[],"source":["## Retreiving the model\n","\n","# Getting the latest version of the artifact\n","artifact = run.use_artifact('{}:latest'.format(run_config['model']))\n","# Downloading the artifact\n","artifact_dir = artifact.download()\n","# Loading the model\n","model_dict = torch.load(os.path.join(artifact_dir, 'Model'))\n","\n","\n","\n","# Loading weights\n","model.load_state_dict(model_dict['model_state_dict'])\n","# Loading optimizer state\n","optimizer.load_state_dict(model_dict['optimizer_state_dict'])"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"2kmyvmCX_Bk_"},"outputs":[],"source":["# Finishing runs\n","wandb.finish()"]}],"metadata":{"accelerator":"GPU","colab":{"provenance":[]},"gpuClass":"standard","kernelspec":{"display_name":"Python 3","name":"python3"},"language_info":{"name":"python"}},"nbformat":4,"nbformat_minor":0}