{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"provenance":[]},"kernelspec":{"name":"python3","display_name":"Python 3"},"language_info":{"name":"python"}},"cells":[{"cell_type":"markdown","source":["**Recitation 0: Checkpointing**\n","\n","We will show you how to checkpoint and load your model :D"],"metadata":{"id":"L8GAh1XaEzCZ"}},{"cell_type":"markdown","source":["**Section 0: Setup**\n","\n","Let's define a quick dummy model, optimizer, and scheduler that we'll be saving and loading :)"],"metadata":{"id":"sGTQZCMESiir"}},{"cell_type":"code","execution_count":null,"metadata":{"id":"JWFlVL5oErGf"},"outputs":[],"source":["!pip install wandb --quiet # Install WandB"]},{"cell_type":"code","source":["import torch\n","import torch.nn as nn\n","import torch.optim as optim\n","import wandb\n","import os\n","\n","device = 'cuda' if torch.cuda.is_available() else 'cpu'\n","print(\"Device: \", device)\n","\n","os.environ['WANDB_API_KEY'] = \"\"#your key here\n","wandb.login()"],"metadata":{"id":"UI51GXlnLdHs","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1735873422013,"user_tz":300,"elapsed":11197,"user":{"displayName":"John Liu","userId":"01307449660175699014"}},"outputId":"9f823d3a-8c02-4201-b2c1-448b10f4113d"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["Device:  cpu\n"]}]},{"cell_type":"code","source":["# A simple submodule\n","class DummySubmodule(nn.Module):\n","    def __init__(self):\n","        super(DummySubmodule, self).__init__()\n","        self.layer = nn.Linear(in_features = 32, out_features = 32)\n","\n","    def forward(self, x):\n","        return self.layer(x)\n","\n","# A simple network\n","class DummyNetwork(nn.Module):\n","\n","    def __init__(self):\n","        super(DummyNetwork, self).__init__()\n","\n","        self.lower_layer = nn.Sequential(\n","            nn.Linear(in_features = 32, out_features = 64),\n","            nn.ReLU(),\n","            nn.Linear(in_features = 32, out_features = 64),\n","            nn.ReLU(),\n","        )\n","        self.upper_layer = nn.Sequential(\n","            nn.Linear(in_features = 64, out_features = 32),\n","            nn.ReLU(),\n","        )\n","\n","        self.module1 = DummySubmodule()\n","\n","    def forward(self, x):\n","        res = self.lower_layer(x)\n","        res = self.upper_layer(res)\n","        res = self.submodule(res)\n","        return res\n","\n","# Declare the model, optimizer, and scheduler\n","model = DummyNetwork().to(device)\n","optimizer = optim.SGD(model.parameters(), lr=0.01)\n","scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)"],"metadata":{"id":"3ZITfNI8Sh6X"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["Let's take a look at some of the information we can checkpoint"],"metadata":{"id":"AP28qWKxXkeN"}},{"cell_type":"code","source":["# Print model's state_dict\n","print(\"==============================================================\")\n","print(\"Model's state_dict:\")\n","for param_tensor in model.state_dict():\n","    print(param_tensor, \"\\t\", model.state_dict()[param_tensor].size())\n","print(\"==============================================================\")\n","# Print optimizer's state_dict\n","print(\"Optimizer's state_dict:\")\n","for var_name in optimizer.state_dict():\n","    print(var_name, \"\\t\", optimizer.state_dict()[var_name])\n","print(\"==============================================================\")\n","# Print scheduler's state_dict\n","print(\"\\nScheduler's state_dict:\")\n","for var_name, value in scheduler.state_dict().items():\n","    print(var_name, \"\\t\", value)\n","print(\"==============================================================\")"],"metadata":{"id":"Trj1bxThUpbI","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1735873330480,"user_tz":300,"elapsed":91,"user":{"displayName":"John Liu","userId":"01307449660175699014"}},"outputId":"a4d0a196-9667-4bb4-a426-4e9e53c53754"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["==============================================================\n","Model's state_dict:\n","lower_layer.0.weight \t torch.Size([64, 32])\n","lower_layer.0.bias \t torch.Size([64])\n","lower_layer.2.weight \t torch.Size([64, 32])\n","lower_layer.2.bias \t torch.Size([64])\n","upper_layer.0.weight \t torch.Size([32, 64])\n","upper_layer.0.bias \t torch.Size([32])\n","module1.layer.weight \t torch.Size([32, 32])\n","module1.layer.bias \t torch.Size([32])\n","==============================================================\n","Optimizer's state_dict:\n","state \t {}\n","param_groups \t [{'lr': 0.01, 'momentum': 0, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'maximize': False, 'foreach': None, 'differentiable': False, 'fused': None, 'initial_lr': 0.01, 'params': [0, 1, 2, 3, 4, 5, 6, 7]}]\n","==============================================================\n","\n","Scheduler's state_dict:\n","step_size \t 10\n","gamma \t 0.1\n","base_lrs \t [0.01]\n","last_epoch \t 0\n","verbose \t False\n","_step_count \t 1\n","_get_lr_called_within_step \t False\n","_last_lr \t [0.01]\n","==============================================================\n"]}]},{"cell_type":"markdown","source":["#######################################################################################################"],"metadata":{"id":"-pA_3jeXbNoz"}},{"cell_type":"markdown","source":["**Section 1: How to save the checkpoint Saving a checkpoint**\n","\n","Checkpointing locally"],"metadata":{"id":"vGmdWx9WSVVM"}},{"cell_type":"code","source":["# let's pretend we're in the middle of our training\n","epoch = 6 # pretend we're in our 6th epoch\n","loss = 0.78 #pretend this is our model's loss at the moment\n","\n","checkpoint_path=f\"<run_name>_{epoch}.pth\"\n","\n","# Saving your states locally with torch.save\n","torch.save({\n","    'model_state_dict': model.state_dict(),   # saving the model state\n","    # if isinstance(model, nn.DataParallel) 'model_state_dict': model.module.state_dict()\n","    'optimizer_state_dict': optimizer.state_dict(),   # saving the optimizer state\n","    'scheduler_state_dict': scheduler.state_dict(),   # saving the scheduler state\n","    'epoch': epoch,\n","    'current_loss': loss\n","    }, checkpoint_path\n",")\n","\n"],"metadata":{"id":"dxyWbA0CSVEg"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["Checkpointing and saving to wandb as an artifact"],"metadata":{"id":"jhCDyi8JX-Dm"}},{"cell_type":"code","source":["# Before the run, you need to have started a run like so....\n","run = wandb.init(\n","    project=\"wandb-quickstart\",\n","    name=\"<run_name>\",\n","    )\n","\n","# ...\n","# ...\n","# ...\n","# Within a training loop (or wherever else you want)....\n","\n","# Option 1:\n","# create artifacts (keeps track of versioning, and is much more organized to work with between collaborators)\n","checkpoint_artifact = wandb.Artifact(\"<run_name>\", type=\"checkpoint\") # You can switch type=\"model if you only want to save a model\"\n","\n","checkpoint_artifact.add_file(checkpoint_path)\n","\n","run.log_artifact(checkpoint_artifact)\n","\n","# Option 2:\n","# directly save the model to wandb\n","wandb.save(checkpoint_path, base_path=os.path.dirname(checkpoint_path))"],"metadata":{"id":"wCZpMYy4YAtM"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["#######################################################################################################"],"metadata":{"id":"ybrAu4ghbdCR"}},{"cell_type":"markdown","source":["**Section 2\\: Loading a checkpoint file into our current model**"],"metadata":{"id":"qW5LJvUnL7d_"}},{"cell_type":"markdown","source":["Downloading a model from wandb"],"metadata":{"id":"EsMCPWbWcEiF"}},{"cell_type":"code","source":["# METHOD 1: Download from wandb Artifact\n","# If you need to re-obtain the run, you can do the following....\n","api = wandb.Api()\n","# information can be obtained from the wandb link adddress as follows:\n","# https://wandb.ai/<USERNAME>/<PROJECT_NAME>/runs/<RUN_ID>?nw=nwuser<USERNAME>\n","run = api.run(\"<USERNAME>/<PROJECT_NAME>/<RUN_ID>\")\n","\n","# To retrieve the artifact....\n","# Get the artifact (choose which version of the model you want)\n","artifact = run.use_artifact('<run_name>:latest')\n","# Downloading the artifact\n","artifact_dir = artifact.download()\n","# Loading the model dict\n","checkpoint_dict = torch.load(os.path.join(artifact_dir, '<run_name>'))\n","\n","\n","# METHOD 2: Download the directly saved file from wandb to Local File\n","checkpoint_file = wandb.restore('<run_name>', run_path=\"<USERNAME>/<PROJECT_NAME>/<RUN_ID>\").name\n","checkpoint_dict = torch.load(checkpoint_file)"],"metadata":{"id":"uFnp3oD5cD0p"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["Loading a .pth checkpoint file from our local directory to our model"],"metadata":{"id":"c0vBwQPkQtMW"}},{"cell_type":"code","source":["# .pth checkpoint file path can also be obtained from a locally saved .pth file. Or, you can use the checkpoint_dict obtained from the prior wandb artifact download :)\n","checkpoint_path = \"/content/<run_name>_6.pth\"\n","checkpoint_dict = torch.load(checkpoint_path)\n","\n","\n","# Loading model weights\n","# if isinstance(model, nn.DataParallel) model.module.load_state_dict(checkpoint_dict['model_state_dict'])\n","model.load_state_dict(checkpoint_dict['model_state_dict'])\n","# Loading optimizer state\n","optimizer.load_state_dict(checkpoint_dict['optimizer_state_dict'])\n","# Loading the scheduler state\n","scheduler.load_state_dict(checkpoint_dict['scheduler_state_dict'])\n","# find the epoch we left off at\n","current_epoch = checkpoint_dict['epoch']\n","# Find any metrics that might be relevant\n","current_loss = checkpoint_dict['current_loss']\n","\n","# Done!!!!!!"],"metadata":{"id":"ZmF238XnPw4j","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1735873481395,"user_tz":300,"elapsed":73,"user":{"displayName":"John Liu","userId":"01307449660175699014"}},"outputId":"9fd6d8ca-6e56-4624-d64f-9427bbae76bc"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stderr","text":["<ipython-input-7-2bf5796c687f>:3: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n","  checkpoint_dict = torch.load(checkpoint_path)\n"]}]},{"cell_type":"code","source":["# If you want to load specific parts of your model (in our case, we can load just the lower layers or just the upper layers)\n","specific_weights = { # Creates dictionary of only desired weights\n","    key: value\n","    for key, value in checkpoint_dict['model_state_dict'].items()\n","    if 'lower_layer' in key\n","}\n","\n","model.load_state_dict(specific_weights, strict=False)"],"metadata":{"id":"NaM3zpYEgybC"},"execution_count":null,"outputs":[]}]}