tutorials/01-basics/06-optim-save-load/main.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "da7db4b0-de84-4120-acd4-d543eda4d137",
   "metadata": {},
   "source": [
    "# OPTIMIZING MODEL PARAMETERS | MODEL SAVE/LOAD\n",
    "\n",
    "Now that we have a model and data it’s time to train, validate and test our model by optimizing its parameters on our data. Training a model is an iterative process; in each iteration the model makes a guess about the output, calculates the error in its guess (loss), collects the derivatives of the error with respect to its parameters, and optimizes these parameters using gradient descent."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b4f57ce-fa19-4b6d-9552-387ae5829a69",
   "metadata": {},
   "source": [
    "## Prerequisite Code\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d14ec220-e6ad-48f0-b674-9d6a0116bf3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets\n",
    "from torchvision.transforms import ToTensor\n",
    "\n",
    "training_data = datasets.FashionMNIST(\n",
    "    root=\"data\",\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=ToTensor()\n",
    ")\n",
    "\n",
    "test_data = datasets.FashionMNIST(\n",
    "    root=\"data\",\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=ToTensor()\n",
    ")\n",
    "\n",
    "train_dataloader = DataLoader(training_data, batch_size=64)\n",
    "test_dataloader = DataLoader(test_data, batch_size=64)\n",
    "\n",
    "class NeuralNetwork(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.linear_relu_stack = nn.Sequential(\n",
    "            nn.Linear(28*28, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 10),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.flatten(x)\n",
    "        logits = self.linear_relu_stack(x)\n",
    "        return logits\n",
    "\n",
    "model = NeuralNetwork()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "471af6e7-1283-4dbb-ab32-133b026e5cc0",
   "metadata": {},
   "source": [
    "## Hyperparameters\n",
    "\n",
    "Hyperparameters are adjustable parameters that let you control the model optimization process. Different hyperparameter values can impact model training and convergence rates.\n",
    "\n",
    "We define the following hyperparameters for training:\n",
    "\n",
    "- **Number of Epochs** - the number times to iterate over the dataset\n",
    "- **Batch Size** - the number of data samples propagated through the network before the parameters are updated\n",
    "- **Learning Rate** - how much to update models parameters at each batch/epoch. Smaller values yield slow learning speed, while large values may result in unpredictable behavior during training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a67c8ba2-e219-4bd5-9bd6-dad5c096c2b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate = 1e-3\n",
    "batch_size = 64\n",
    "epochs = 5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9500d5a-96c6-42a7-b080-69755559fabf",
   "metadata": {},
   "source": [
    "## Optimization Loop\n",
    "\n",
    "Once we set our hyperparameters, we can then train and optimize our model with an optimization loop. Each iteration of the optimization loop is called an epoch.\n",
    "\n",
    "Each epoch consists of two main parts:\n",
    "- **The Train Loop** - iterate over the training dataset and try to converge to optimal parameters.\n",
    "- **The Validation/Test Loop** - iterate over the test dataset to check if model performance is improving."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96361b23-66d1-4610-bb3f-88315fd5d30d",
   "metadata": {},
   "source": [
    "### Loss Function and Optimizer\n",
    "\n",
    "Optimization is the process of adjusting model parameters to reduce model error in each training step. Optimization algorithms define how this process is performed (in this example we use Stochastic Gradient Descent)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a0592bb6-afae-4613-94a1-e2c1759a114a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the loss function\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "\n",
    "# Opimizer\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "763ac490-b571-431d-b185-40f43b6faa55",
   "metadata": {},
   "source": [
    "Inside the training loop, optimization happens in three steps:\n",
    "- Call `optimizer.zero_grad(`) to reset the gradients of model parameters. Gradients by default add up; to prevent double-counting, we explicitly zero them at each iteration.\n",
    "- Backpropagate the prediction loss with a call to `loss.backward()`. PyTorch deposits the gradients of the loss w.r.t. each parameter.\n",
    "- Once we have our gradients, we call `optimizer.step()` to adjust the parameters by the gradients collected in the backward pass."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "de3ecc25-e4e4-4adb-949b-0e5bc70cad28",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_loop(dataloader, model, loss_fn, optimizer):\n",
    "    size = len(dataloader.dataset)\n",
    "    # Set the model to training mode - important for batch normalization and dropout layers\n",
    "    # Unnecessary in this situation but added for best practices\n",
    "    model.train()\n",
    "    for batch, (X, y) in enumerate(dataloader):\n",
    "        # Compute prediction and loss\n",
    "        pred = model(X)\n",
    "        loss = loss_fn(pred, y)\n",
    "\n",
    "        # Backpropagation\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        if batch % 100 == 0:\n",
    "            loss, current = loss.item(), (batch + 1) * len(X)\n",
    "            print(f\"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]\")\n",
    "\n",
    "\n",
    "def test_loop(dataloader, model, loss_fn):\n",
    "    # Set the model to evaluation mode - important for batch normalization and dropout layers\n",
    "    # Unnecessary in this situation but added for best practices\n",
    "    model.eval()\n",
    "    size = len(dataloader.dataset)\n",
    "    num_batches = len(dataloader)\n",
    "    test_loss, correct = 0, 0\n",
    "\n",
    "    # Evaluating the model with torch.no_grad() ensures that no gradients are computed during test mode\n",
    "    # also serves to reduce unnecessary gradient computations and memory usage for tensors with requires_grad=True\n",
    "    with torch.no_grad():\n",
    "        for X, y in dataloader:\n",
    "            pred = model(X)\n",
    "            test_loss += loss_fn(pred, y).item()\n",
    "            correct += (pred.argmax(1) == y).type(torch.float).sum().item()\n",
    "\n",
    "    test_loss /= num_batches\n",
    "    correct /= size\n",
    "    print(f\"Test Error: \\n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "73e45919-941b-4ce6-a1f2-d920345b760a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1\n",
      "-------------------------------\n",
      "loss: 2.308445  [   64/60000]\n",
      "loss: 2.298456  [ 6464/60000]\n",
      "loss: 2.276001  [12864/60000]\n",
      "loss: 2.265796  [19264/60000]\n",
      "loss: 2.244822  [25664/60000]\n",
      "loss: 2.219735  [32064/60000]\n",
      "loss: 2.228179  [38464/60000]\n",
      "loss: 2.189170  [44864/60000]\n",
      "loss: 2.177993  [51264/60000]\n",
      "loss: 2.154140  [57664/60000]\n",
      "Test Error: \n",
      " Accuracy: 42.2%, Avg loss: 2.146824 \n",
      "\n",
      "Epoch 2\n",
      "-------------------------------\n",
      "loss: 2.156378  [   64/60000]\n",
      "loss: 2.151325  [ 6464/60000]\n",
      "loss: 2.082079  [12864/60000]\n",
      "loss: 2.095107  [19264/60000]\n",
      "loss: 2.047306  [25664/60000]\n",
      "loss: 1.982480  [32064/60000]\n",
      "loss: 2.020417  [38464/60000]\n",
      "loss: 1.930593  [44864/60000]\n",
      "loss: 1.933225  [51264/60000]\n",
      "loss: 1.867908  [57664/60000]\n",
      "Test Error: \n",
      " Accuracy: 51.1%, Avg loss: 1.862668 \n",
      "\n",
      "Epoch 3\n",
      "-------------------------------\n",
      "loss: 1.898650  [   64/60000]\n",
      "loss: 1.872178  [ 6464/60000]\n",
      "loss: 1.739482  [12864/60000]\n",
      "loss: 1.782994  [19264/60000]\n",
      "loss: 1.681339  [25664/60000]\n",
      "loss: 1.629598  [32064/60000]\n",
      "loss: 1.671012  [38464/60000]\n",
      "loss: 1.562859  [44864/60000]\n",
      "loss: 1.589980  [51264/60000]\n",
      "loss: 1.497004  [57664/60000]\n",
      "Test Error: \n",
      " Accuracy: 60.2%, Avg loss: 1.505618 \n",
      "\n",
      "Epoch 4\n",
      "-------------------------------\n",
      "loss: 1.572734  [   64/60000]\n",
      "loss: 1.544055  [ 6464/60000]\n",
      "loss: 1.384945  [12864/60000]\n",
      "loss: 1.459769  [19264/60000]\n",
      "loss: 1.349101  [25664/60000]\n",
      "loss: 1.338249  [32064/60000]\n",
      "loss: 1.374898  [38464/60000]\n",
      "loss: 1.288753  [44864/60000]\n",
      "loss: 1.319918  [51264/60000]\n",
      "loss: 1.236561  [57664/60000]\n",
      "Test Error: \n",
      " Accuracy: 62.9%, Avg loss: 1.251954 \n",
      "\n",
      "Epoch 5\n",
      "-------------------------------\n",
      "loss: 1.325233  [   64/60000]\n",
      "loss: 1.313260  [ 6464/60000]\n",
      "loss: 1.144201  [12864/60000]\n",
      "loss: 1.247604  [19264/60000]\n",
      "loss: 1.129655  [25664/60000]\n",
      "loss: 1.147383  [32064/60000]\n",
      "loss: 1.189313  [38464/60000]\n",
      "loss: 1.115911  [44864/60000]\n",
      "loss: 1.146612  [51264/60000]\n",
      "loss: 1.077964  [57664/60000]\n",
      "Test Error: \n",
      " Accuracy: 64.5%, Avg loss: 1.090530 \n",
      "\n",
      "Epoch 6\n",
      "-------------------------------\n",
      "loss: 1.157466  [   64/60000]\n",
      "loss: 1.162811  [ 6464/60000]\n",
      "loss: 0.981731  [12864/60000]\n",
      "loss: 1.110013  [19264/60000]\n",
      "loss: 0.990372  [25664/60000]\n",
      "loss: 1.016198  [32064/60000]\n",
      "loss: 1.070604  [38464/60000]\n",
      "loss: 1.002520  [44864/60000]\n",
      "loss: 1.031578  [51264/60000]\n",
      "loss: 0.975071  [57664/60000]\n",
      "Test Error: \n",
      " Accuracy: 65.5%, Avg loss: 0.984124 \n",
      "\n",
      "Epoch 7\n",
      "-------------------------------\n",
      "loss: 1.039809  [   64/60000]\n",
      "loss: 1.063433  [ 6464/60000]\n",
      "loss: 0.868616  [12864/60000]\n",
      "loss: 1.016345  [19264/60000]\n",
      "loss: 0.900921  [25664/60000]\n",
      "loss: 0.922687  [32064/60000]\n",
      "loss: 0.991155  [38464/60000]\n",
      "loss: 0.927610  [44864/60000]\n",
      "loss: 0.952159  [51264/60000]\n",
      "loss: 0.905622  [57664/60000]\n",
      "Test Error: \n",
      " Accuracy: 66.6%, Avg loss: 0.911320 \n",
      "\n",
      "Epoch 8\n",
      "-------------------------------\n",
      "loss: 0.953631  [   64/60000]\n",
      "loss: 0.994393  [ 6464/60000]\n",
      "loss: 0.787500  [12864/60000]\n",
      "loss: 0.950229  [19264/60000]\n",
      "loss: 0.841003  [25664/60000]\n",
      "loss: 0.854559  [32064/60000]\n",
      "loss: 0.935056  [38464/60000]\n",
      "loss: 0.877088  [44864/60000]\n",
      "loss: 0.895778  [51264/60000]\n",
      "loss: 0.856210  [57664/60000]\n",
      "Test Error: \n",
      " Accuracy: 67.7%, Avg loss: 0.859335 \n",
      "\n",
      "Epoch 9\n",
      "-------------------------------\n",
      "loss: 0.888271  [   64/60000]\n",
      "loss: 0.943319  [ 6464/60000]\n",
      "loss: 0.727435  [12864/60000]\n",
      "loss: 0.902059  [19264/60000]\n",
      "loss: 0.798541  [25664/60000]\n",
      "loss: 0.803932  [32064/60000]\n",
      "loss: 0.892868  [38464/60000]\n",
      "loss: 0.841733  [44864/60000]\n",
      "loss: 0.854251  [51264/60000]\n",
      "loss: 0.819232  [57664/60000]\n",
      "Test Error: \n",
      " Accuracy: 68.9%, Avg loss: 0.820435 \n",
      "\n",
      "Epoch 10\n",
      "-------------------------------\n",
      "loss: 0.836433  [   64/60000]\n",
      "loss: 0.903223  [ 6464/60000]\n",
      "loss: 0.681219  [12864/60000]\n",
      "loss: 0.865816  [19264/60000]\n",
      "loss: 0.766572  [25664/60000]\n",
      "loss: 0.765379  [32064/60000]\n",
      "loss: 0.858930  [38464/60000]\n",
      "loss: 0.815500  [44864/60000]\n",
      "loss: 0.822285  [51264/60000]\n",
      "loss: 0.790205  [57664/60000]\n",
      "Test Error: \n",
      " Accuracy: 70.0%, Avg loss: 0.789832 \n",
      "\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "loss_fn = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)\n",
    "\n",
    "epochs = 10\n",
    "for t in range(epochs):\n",
    "    print(f\"Epoch {t+1}\\n-------------------------------\")\n",
    "    train_loop(train_dataloader, model, loss_fn, optimizer)\n",
    "    test_loop(test_dataloader, model, loss_fn)\n",
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e418e2fd-ef11-4855-b9da-6b191bcb2f47",
   "metadata": {},
   "source": [
    "## Save and Load the Model\n",
    "\n",
    "PyTorch models store the learned parameters in an internal state dictionary, called `state_dict`. These can be persisted via the `torch.save` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6883276f-b104-45e4-b531-fddc9f694887",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), \"model.pth\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45e384d7-996b-4d17-9485-903312777b50",
   "metadata": {},
   "source": [
    "To load model weights, you need to create an **instance of the same model first**, and then load the parameters using `load_state_dict()` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "95d2618f-4d8e-4fc8-9641-d12eb249a690",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net = NeuralNetwork()\n",
    "net.load_state_dict(torch.load(\"model.pth\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "022ae63e-837f-4ad7-a9f6-9530ab165e87",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "NeuralNetwork(\n",
       "  (flatten): Flatten(start_dim=1, end_dim=-1)\n",
       "  (linear_relu_stack): Sequential(\n",
       "    (0): Linear(in_features=784, out_features=512, bias=True)\n",
       "    (1): ReLU()\n",
       "    (2): Linear(in_features=512, out_features=512, bias=True)\n",
       "    (3): ReLU()\n",
       "    (4): Linear(in_features=512, out_features=10, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net.eval() # sets the model to evaluation mode"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b01e5fd8-8f9a-42de-b5e3-15ff5041f3a3",
   "metadata": {},
   "source": [
    "## Saving and Loading Models with Shapes\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6846e0ff-f4f2-4ff7-b888-233b227e418b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# saving\n",
    "torch.save(net, \"model_with_shapes.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "da54d613-1d3b-48cf-9e08-b88aa08eddd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# loading\n",
    "model = torch.load('model_with_shapes.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2a136162-2def-467e-bbc0-bda329f1c080",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "NeuralNetwork(\n",
       "  (flatten): Flatten(start_dim=1, end_dim=-1)\n",
       "  (linear_relu_stack): Sequential(\n",
       "    (0): Linear(in_features=784, out_features=512, bias=True)\n",
       "    (1): ReLU()\n",
       "    (2): Linear(in_features=512, out_features=512, bias=True)\n",
       "    (3): ReLU()\n",
       "    (4): Linear(in_features=512, out_features=10, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1200a7e1-20bb-4161-8c39-cef4187dd6df",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}