Part 6 - Saving and Loading Models.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Saving and Loading Models\n",
    "\n",
    "In this notebook, I'll show you how to save and load models with PyTorch. This is important because you'll often want to load previously trained models to use in making predictions or to continue training on new data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch import optim\n",
    "import torch.nn.functional as F\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "import helper\n",
    "import fc_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a transform to normalize the data\n",
    "transform = transforms.Compose([transforms.ToTensor(),\n",
    "                                transforms.Normalize((0.5,), (0.5,))])\n",
    "# Download and load the training data\n",
    "trainset = datasets.FashionMNIST('~/.pytorch/F_MNIST_data/', download=True, train=True, transform=transform)\n",
    "trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)\n",
    "\n",
    "# Download and load the test data\n",
    "testset = datasets.FashionMNIST('~/.pytorch/F_MNIST_data/', download=True, train=False, transform=transform)\n",
    "testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we can see one of the images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAc8AAAHPCAYAAAA1eFErAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAABYlAAAWJQFJUiTwAAAP2ElEQVR4nO3dy46c12GF0VNVXdW8mSYBE7okloQgnvqSAM4kkzizIICRmZF3C/wmmcSZxA6cAHFkS7KkSKIMkxApkRSbzbpl4AeI99kBG41ea354qqqr+fU/2ovj8TgAgD/e8qJfAABcNuIJACHxBICQeAJASDwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASB0MnvwR3/9PXMs/NGuXbs2ffYffvzj6u6bN29On71750519yeffnZhd7d+/+DB9Nn3P3i/uvu/fvWr6vxVtFgsqvNXdWHrn3/2n1MfnCdPAAiJJwCExBMAQuIJACHxBICQeAJASDwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASA0vefJq7dcdn/rvP3W29Nnv/OdP6/uvvete9Nnz89fVHe///4H02d3u1119/n5+fTZ5ar7eS8X3fnN6Wb67Dtvv1Pd/e0//fb02S8efVHd/dln8xusn//ud9Xdjau6x3lRPHkCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQibJQjdu3KjOf++7350++87b85NiY4yxKCaqzs6eV3c/fvx4+mw7xXb79jemzz59+qy6e7efnzQ73cxPgo0xxqY8v1ytps8+e9Z9bqtiju2N19+o7n77rbemzz5/3v2e3P/88+mz//3uu9XdL1++rM5fNZ48ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJACHxBIDQldzz/MH3v38hZ8cY4/cPHkyfbTYxxxhjt5vflmy2QMcYY7FcTJ9dL7qv6XI5f/7WrZvV3evNevrs6ea0untfbImOMcbzs7Pps4sx//MeY4zdfj999vHjR9XdYzH/2k+KDdQxxrj3rXvTZ//+7/6suvuX//HL6bMfffxxdfdl5MkTAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AELqSk2Tn5y+nz37xqJs7Oh4O02d3u/mZpjHGWBZzSatlN7VUrDyNQ/GZjTHGixfP5g83L3yMsSo+8/Pz8+ru7W5bnW++q2NZTtg1n3s5n3c8zr/vbTH7N8YYh+Nx/uyh+/+h+a5eRZ48ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJACHxBIDQldzz/Ph/Pp4++5d/8YPq7sePv5w+e+1at7e33c7vmO733VbgKOYZl+0+Y3H2UL7v5nNbr9fV3YvmQx9jFNOS5eExDvv5Tc1ygrX73Nv912I3t90S/eC3v63OXzWePAEgJJ4AEBJPAAiJJwCExBMAQuIJACHxBICQeAJASDwBICSeABASTwAIiScAhMQTAELiCQChKzlJ9vz58+LsWXV3s1i03XaTQ5vNZvrs+cv5ObMxqkWyeg7tcJift7p181Z1dzOIttt17/vlsfuZLZfzf1sv2l2w4n+mdsrtWHxfzs+7z/zevXvTZ3/x77+o7ibjyRMAQuIJACHxBICQeAJASDwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACF3JPc/Gu7/+dXX+r374w+mzH338UXX3GNemTy4X3d9Zh+P8RmK7Dbko1kRfbrt9xkOxRbrddfutrdXJav7scv7sGKMagG33X1+8eDF99tatbv/1wYMH02c/+fTT6m4ynjwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAIZNkoffef686/7c/+pvps6ebTXX3/jA/1dROTDWjYodxrO4+OZn/mjeTYmOMsSim3NbF6x5jjMOx+9yqIbhuRW5st9vps4fD/PzdGGNsit+z9nf0337+8+o8r44nTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgZM8ztNvtqvP3738+ffb1116v7v7s/mfTZ1fL7u+sZlvyeCh3KRfz45KL8n03U6Tdu+7e9xjlBmu5qXksvi/LYkN1jDHW6/X02SdPn1Z3f/nll9V5Xh1PngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJACHxBICQSbJX7F/+9WfTZ//xJz+p7n729dfTZ58/nz87Rj9RVSnmrdpdsEUxj7Vs59Au0H7bTfcdi+/Lvpywu379+vTZDz75oLqby+Py/nYCwAURTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkDInucr9vDhw+mzH374YXX33bt3p89ut9vq7jEW0yePx4vbAj00W6ClVb3nOf+ZjzHGsXjvi0139+G8+Jkf9tXdy9X85/7kydPqbi4PT54AEBJPAAiJJwCExBMAQuIJACHxBICQeAJASDwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkEmyS+SLR4+q82+++eb02d1uV919dvZi+mw7h3YcxbTWvptDa2a9lqtVdfeinSQrPrfVsfu7fL+fnxVbLrq7j4f5n/mTp0+quxuLxcVN0F1FnjwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJA9z0vkUbnneSh2Ck9Ouq/KatX8ndbd3WxDNtuOY3Sf+fHQ7TOOZfe3cXP7/tBtQzafe7NDOsYYy2W3o8rV4MkTAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEDJJFlosupmo43F+Lun8/Ly6u5nHWlQDVWNsd7v5w93CVGVRznoti5/3oZz1Gof5KbYxxmi+6vtyym1XzMiN4jMfo/w9Kf9/4PLw5AkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJACHxBICQeAJASDwBICSeABASTwAIiScAhOx5XiLb7bY632yJHstRzWWxc7jdde97tZr/mq+W5X5rscl5PHZ7nKPcYK3OX+AGa+t4tOfJ/82TJwCExBMAQuIJACHxBICQeAJASDwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASAkngAQsucZajYxW/t9t+9Y7XmW73uxmP877VBsYo4xxnrdbCx2+4y73fzZ/WF+V3KMMVbL7m/j5mfWzlouis+93Z69wF9xLhFPngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJACHxBICQSbJLZLlaVeebqaV2FmxRbVR1dy+X85/bspz1Oj9/OX32UE6SdZ/5GCfL5nz3XV0Udx+75b5xPHafO1eDJ08ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJACHxBICQeAJASDwBIGTPM9RuJB6LUc1leXfz2tuNw2Wzz9jNeVabnOv1xf2KHMs9z1Huvzbfl3YHtfmu78vvy6H4wrX/P3B5ePIEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJACHxBICQeAJASDwBICSeABASTwAIiScAhEyShdrBoWYt6eSk/XHN374q563Ozl4Up8uNqUL7vptprmYaa4wxlovub+Pm+naaa1V815+fnVV331rdmj577dq16u4XL5rfE14lT54AEBJPAAiJJwCExBMAQuIJACHxBICQeAJASDwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQMieZ+jiliXH2Gw21fnr169Pn/3666+ruxvNJuYYY5yczG9ytnc3juWeZ/ttPRz202fbHdT1yXr+cPmxLVfzP/P2fTcucmv4KvLkCQAh8QSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiZJAv1M1Hzbty4UZ1//vxs+mz7vo/Hw/TZduapOb9cdnc3k2aLcmRqsSjPF6+9/b6s1/OTZM2k2BhjHPbz39ULnbC7sJuvJk+eABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkDInucl8idvvln+C/OLfxe5a7nf7y/s7mW5iblczp9v9zhHuQe6Kj637XZX3b1ez7/2dv+1+ZltNpvq7sZFbg1fRZ48ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJACGTZJfIG2+8UZ0/Pz+fPtvMeo0xxunp6fTZ3a6bt2omqtr33ZxvF8n6882cWnf3fn+YPnu6mf+ujdHNqd29c6e6+/79+9V5Xh1PngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJACHxBICQeAJA6ErueTY7hcfj8f/xlWSePn1Wnd9s1tNnm8+sP9/dXW1qLi/wfbejmOXndjjMb2o2G6qt9br7b61539+8/c3qbi4PT54AEBJPAAiJJwCExBMAQuIJACHxBICQeAJASDwBICSeABASTwAIiScAhMQTAELiCQAh8QSA0NWcJCvOtoNkr7322vTZ27e/Ud292+2mz7ZLbPv9/MxTMxH1h/PFiz/uq7v3++J8+aEfDt1rb77t7XRfM2m2WHTPBIvF/Gu/fv16dTeXhydPAAiJJwCExBMAQuIJACHxBICQeAJASDwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASB0Jfc8203OxsOHD6fPVtuQY4zNejN99vTaaXX3jWLn8ObNG9XdN2/cnD67WDTrr6Maj23vvlG87z+Y/01p9lvHGGO9nv+vqdmtHaObUf3Ne7+p7uby8OQJACHxBICQeAJASDwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACF3NSbJmc6h0OMxPNf3TT39a3b1arabP3r1zp7r79u3b02fPX76s7r51c36a6/S0m2J7Wbz2r776qrp7uez+Nm7Ob7fdLFizxnYof7+fPHkyffbs7Ky6u9FO2F3k/4uXkSdPAAiJJwCExBMAQuIJACHxBICQeAJASDwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASC0sOEGABlPngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJAKH/BZCVqBP9JQi+AAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 231,
       "width": 231
      },
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "image, label = next(iter(trainloader))\n",
    "helper.imshow(image[0,:]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train a network\n",
    "\n",
    "To make things more concise here, I moved the model architecture and training code from the last part to a file called `fc_model`. Importing this, we can easily create a fully-connected network with `fc_model.Network`, and train the network using `fc_model.train`. I'll use this model (once it's trained) to demonstrate how we can save and load models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the network, define the criterion and optimizer\n",
    "\n",
    "model = fc_model.Network(784, 10, [512, 256, 128])\n",
    "criterion = nn.NLLLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 1/2..  Training Loss: 1.653..  Test Loss: 0.967..  Test Accuracy: 0.649\n",
      "Epoch: 1/2..  Training Loss: 1.015..  Test Loss: 0.718..  Test Accuracy: 0.726\n",
      "Epoch: 1/2..  Training Loss: 0.859..  Test Loss: 0.665..  Test Accuracy: 0.737\n",
      "Epoch: 1/2..  Training Loss: 0.754..  Test Loss: 0.636..  Test Accuracy: 0.765\n",
      "Epoch: 1/2..  Training Loss: 0.727..  Test Loss: 0.602..  Test Accuracy: 0.767\n",
      "Epoch: 1/2..  Training Loss: 0.714..  Test Loss: 0.613..  Test Accuracy: 0.761\n",
      "Epoch: 1/2..  Training Loss: 0.674..  Test Loss: 0.581..  Test Accuracy: 0.781\n",
      "Epoch: 1/2..  Training Loss: 0.681..  Test Loss: 0.591..  Test Accuracy: 0.777\n",
      "Epoch: 1/2..  Training Loss: 0.666..  Test Loss: 0.556..  Test Accuracy: 0.796\n",
      "Epoch: 1/2..  Training Loss: 0.659..  Test Loss: 0.570..  Test Accuracy: 0.781\n",
      "Epoch: 1/2..  Training Loss: 0.603..  Test Loss: 0.534..  Test Accuracy: 0.805\n",
      "Epoch: 1/2..  Training Loss: 0.619..  Test Loss: 0.534..  Test Accuracy: 0.794\n",
      "Epoch: 1/2..  Training Loss: 0.639..  Test Loss: 0.535..  Test Accuracy: 0.802\n",
      "Epoch: 1/2..  Training Loss: 0.582..  Test Loss: 0.516..  Test Accuracy: 0.811\n",
      "Epoch: 1/2..  Training Loss: 0.613..  Test Loss: 0.502..  Test Accuracy: 0.814\n",
      "Epoch: 1/2..  Training Loss: 0.581..  Test Loss: 0.506..  Test Accuracy: 0.813\n",
      "Epoch: 1/2..  Training Loss: 0.616..  Test Loss: 0.502..  Test Accuracy: 0.815\n",
      "Epoch: 1/2..  Training Loss: 0.572..  Test Loss: 0.500..  Test Accuracy: 0.810\n",
      "Epoch: 1/2..  Training Loss: 0.584..  Test Loss: 0.495..  Test Accuracy: 0.816\n",
      "Epoch: 1/2..  Training Loss: 0.569..  Test Loss: 0.489..  Test Accuracy: 0.821\n",
      "Epoch: 1/2..  Training Loss: 0.571..  Test Loss: 0.493..  Test Accuracy: 0.821\n",
      "Epoch: 1/2..  Training Loss: 0.577..  Test Loss: 0.506..  Test Accuracy: 0.820\n",
      "Epoch: 1/2..  Training Loss: 0.581..  Test Loss: 0.509..  Test Accuracy: 0.815\n",
      "Epoch: 2/2..  Training Loss: 0.543..  Test Loss: 0.490..  Test Accuracy: 0.819\n",
      "Epoch: 2/2..  Training Loss: 0.564..  Test Loss: 0.471..  Test Accuracy: 0.825\n",
      "Epoch: 2/2..  Training Loss: 0.585..  Test Loss: 0.477..  Test Accuracy: 0.826\n",
      "Epoch: 2/2..  Training Loss: 0.539..  Test Loss: 0.472..  Test Accuracy: 0.829\n",
      "Epoch: 2/2..  Training Loss: 0.518..  Test Loss: 0.465..  Test Accuracy: 0.828\n",
      "Epoch: 2/2..  Training Loss: 0.514..  Test Loss: 0.460..  Test Accuracy: 0.831\n",
      "Epoch: 2/2..  Training Loss: 0.536..  Test Loss: 0.474..  Test Accuracy: 0.825\n",
      "Epoch: 2/2..  Training Loss: 0.561..  Test Loss: 0.458..  Test Accuracy: 0.826\n",
      "Epoch: 2/2..  Training Loss: 0.509..  Test Loss: 0.467..  Test Accuracy: 0.830\n",
      "Epoch: 2/2..  Training Loss: 0.499..  Test Loss: 0.449..  Test Accuracy: 0.836\n",
      "Epoch: 2/2..  Training Loss: 0.572..  Test Loss: 0.460..  Test Accuracy: 0.838\n",
      "Epoch: 2/2..  Training Loss: 0.526..  Test Loss: 0.465..  Test Accuracy: 0.832\n",
      "Epoch: 2/2..  Training Loss: 0.541..  Test Loss: 0.467..  Test Accuracy: 0.827\n",
      "Epoch: 2/2..  Training Loss: 0.550..  Test Loss: 0.476..  Test Accuracy: 0.824\n",
      "Epoch: 2/2..  Training Loss: 0.520..  Test Loss: 0.450..  Test Accuracy: 0.837\n",
      "Epoch: 2/2..  Training Loss: 0.517..  Test Loss: 0.446..  Test Accuracy: 0.841\n",
      "Epoch: 2/2..  Training Loss: 0.512..  Test Loss: 0.441..  Test Accuracy: 0.838\n",
      "Epoch: 2/2..  Training Loss: 0.506..  Test Loss: 0.458..  Test Accuracy: 0.833\n",
      "Epoch: 2/2..  Training Loss: 0.560..  Test Loss: 0.438..  Test Accuracy: 0.837\n",
      "Epoch: 2/2..  Training Loss: 0.481..  Test Loss: 0.446..  Test Accuracy: 0.838\n",
      "Epoch: 2/2..  Training Loss: 0.511..  Test Loss: 0.458..  Test Accuracy: 0.831\n",
      "Epoch: 2/2..  Training Loss: 0.504..  Test Loss: 0.460..  Test Accuracy: 0.827\n",
      "Epoch: 2/2..  Training Loss: 0.546..  Test Loss: 0.445..  Test Accuracy: 0.836\n"
     ]
    }
   ],
   "source": [
    "fc_model.train(model, trainloader, testloader, criterion, optimizer, epochs=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Saving and loading networks\n",
    "\n",
    "As you can imagine, it's impractical to train a network every time you need to use it. Instead, we can save trained networks then load them later to train more or use them for predictions.\n",
    "\n",
    "The parameters for PyTorch networks are stored in a model's `state_dict`. We can see the state dict contains the weight and bias matrices for each of our layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Our model: \n",
      "\n",
      " Network(\n",
      "  (hidden_layers): ModuleList(\n",
      "    (0): Linear(in_features=784, out_features=512, bias=True)\n",
      "    (1): Linear(in_features=512, out_features=256, bias=True)\n",
      "    (2): Linear(in_features=256, out_features=128, bias=True)\n",
      "  )\n",
      "  (output): Linear(in_features=128, out_features=10, bias=True)\n",
      "  (dropout): Dropout(p=0.5, inplace=False)\n",
      ") \n",
      "\n",
      "The state dict keys: \n",
      "\n",
      " odict_keys(['hidden_layers.0.weight', 'hidden_layers.0.bias', 'hidden_layers.1.weight', 'hidden_layers.1.bias', 'hidden_layers.2.weight', 'hidden_layers.2.bias', 'output.weight', 'output.bias'])\n"
     ]
    }
   ],
   "source": [
    "print(\"Our model: \\n\\n\", model, '\\n')\n",
    "print(\"The state dict keys: \\n\\n\", model.state_dict().keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The simplest thing to do is simply save the state dict with `torch.save`. For example, we can save it to a file `'checkpoint.pth'`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), 'checkpoint.pth')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we can load the state dict with `torch.load`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "odict_keys(['hidden_layers.0.weight', 'hidden_layers.0.bias', 'hidden_layers.1.weight', 'hidden_layers.1.bias', 'hidden_layers.2.weight', 'hidden_layers.2.bias', 'output.weight', 'output.bias'])\n"
     ]
    }
   ],
   "source": [
    "state_dict = torch.load('checkpoint.pth')\n",
    "print(state_dict.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And to load the state dict in to the network, you do `model.load_state_dict(state_dict)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.load_state_dict(state_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Seems pretty straightforward, but as usual it's a bit more complicated. Loading the state dict works only if the model architecture is exactly the same as the checkpoint architecture. If I create a model with a different architecture, this fails."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "Error(s) in loading state_dict for Network:\n\tsize mismatch for hidden_layers.0.weight: copying a param with shape torch.Size([512, 784]) from checkpoint, the shape in current model is torch.Size([400, 784]).\n\tsize mismatch for hidden_layers.0.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([400]).\n\tsize mismatch for hidden_layers.1.weight: copying a param with shape torch.Size([256, 512]) from checkpoint, the shape in current model is torch.Size([200, 400]).\n\tsize mismatch for hidden_layers.1.bias: copying a param with shape torch.Size([256]) from checkpoint, the shape in current model is torch.Size([200]).\n\tsize mismatch for hidden_layers.2.weight: copying a param with shape torch.Size([128, 256]) from checkpoint, the shape in current model is torch.Size([100, 200]).\n\tsize mismatch for hidden_layers.2.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([100]).\n\tsize mismatch for output.weight: copying a param with shape torch.Size([10, 128]) from checkpoint, the shape in current model is torch.Size([10, 100]).",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Input \u001b[0;32mIn [10]\u001b[0m, in \u001b[0;36m<cell line: 4>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      2\u001b[0m model \u001b[38;5;241m=\u001b[39m fc_model\u001b[38;5;241m.\u001b[39mNetwork(\u001b[38;5;241m784\u001b[39m, \u001b[38;5;241m10\u001b[39m, [\u001b[38;5;241m400\u001b[39m, \u001b[38;5;241m200\u001b[39m, \u001b[38;5;241m100\u001b[39m])\n\u001b[1;32m      3\u001b[0m \u001b[38;5;66;03m# This will throw an error because the tensor sizes are wrong!\u001b[39;00m\n\u001b[0;32m----> 4\u001b[0m \u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mload_state_dict\u001b[49m\u001b[43m(\u001b[49m\u001b[43mstate_dict\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/torchy/lib/python3.8/site-packages/torch/nn/modules/module.py:1497\u001b[0m, in \u001b[0;36mModule.load_state_dict\u001b[0;34m(self, state_dict, strict)\u001b[0m\n\u001b[1;32m   1492\u001b[0m         error_msgs\u001b[38;5;241m.\u001b[39minsert(\n\u001b[1;32m   1493\u001b[0m             \u001b[38;5;241m0\u001b[39m, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mMissing key(s) in state_dict: \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m. \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mformat(\n\u001b[1;32m   1494\u001b[0m                 \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m, \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mjoin(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mformat(k) \u001b[38;5;28;01mfor\u001b[39;00m k \u001b[38;5;129;01min\u001b[39;00m missing_keys)))\n\u001b[1;32m   1496\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(error_msgs) \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[0;32m-> 1497\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mError(s) in loading state_dict for \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m:\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\t\u001b[39;00m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mformat(\n\u001b[1;32m   1498\u001b[0m                        \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\t\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mjoin(error_msgs)))\n\u001b[1;32m   1499\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m _IncompatibleKeys(missing_keys, unexpected_keys)\n",
      "\u001b[0;31mRuntimeError\u001b[0m: Error(s) in loading state_dict for Network:\n\tsize mismatch for hidden_layers.0.weight: copying a param with shape torch.Size([512, 784]) from checkpoint, the shape in current model is torch.Size([400, 784]).\n\tsize mismatch for hidden_layers.0.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([400]).\n\tsize mismatch for hidden_layers.1.weight: copying a param with shape torch.Size([256, 512]) from checkpoint, the shape in current model is torch.Size([200, 400]).\n\tsize mismatch for hidden_layers.1.bias: copying a param with shape torch.Size([256]) from checkpoint, the shape in current model is torch.Size([200]).\n\tsize mismatch for hidden_layers.2.weight: copying a param with shape torch.Size([128, 256]) from checkpoint, the shape in current model is torch.Size([100, 200]).\n\tsize mismatch for hidden_layers.2.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([100]).\n\tsize mismatch for output.weight: copying a param with shape torch.Size([10, 128]) from checkpoint, the shape in current model is torch.Size([10, 100])."
     ]
    }
   ],
   "source": [
    "# Try this\n",
    "model = fc_model.Network(784, 10, [400, 200, 100])\n",
    "# This will throw an error because the tensor sizes are wrong!\n",
    "model.load_state_dict(state_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This means we need to rebuild the model exactly as it was when trained. Information about the model architecture needs to be saved in the checkpoint, along with the state dict. To do this, you build a dictionary with all the information you need to compeletely rebuild the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = {'input_size': 784,\n",
    "              'output_size': 10,\n",
    "              'hidden_layers': [each.out_features for each in model.hidden_layers],\n",
    "              'state_dict': model.state_dict()}\n",
    "\n",
    "torch.save(checkpoint, 'checkpoint.pth')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the checkpoint has all the necessary information to rebuild the trained model. You can easily make that a function if you want. Similarly, we can write a function to load checkpoints. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_checkpoint(filepath):\n",
    "    checkpoint = torch.load(filepath)\n",
    "    model = fc_model.Network(checkpoint['input_size'],\n",
    "                             checkpoint['output_size'],\n",
    "                             checkpoint['hidden_layers'])\n",
    "    model.load_state_dict(checkpoint['state_dict'])\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Network(\n",
      "  (hidden_layers): ModuleList(\n",
      "    (0): Linear(in_features=784, out_features=400, bias=True)\n",
      "    (1): Linear(in_features=400, out_features=200, bias=True)\n",
      "    (2): Linear(in_features=200, out_features=100, bias=True)\n",
      "  )\n",
      "  (output): Linear(in_features=100, out_features=10, bias=True)\n",
      "  (dropout): Dropout(p=0.5, inplace=False)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "model = load_checkpoint('checkpoint.pth')\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}