Part 6 - Saving and Loading Models.ipynb

{
 "metadata": {
  "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.0"
  },
  "orig_nbformat": 2,
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.8.0 64-bit (conda)",
   "metadata": {
    "interpreter": {
     "hash": "03877239ae18a66c3d863e76fda72b4917823ea0b4359fa63eeccf9c06b3cab7"
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2,
 "cells": [
  {
   "source": [
    "# Saving and Loading Models\n",
    "We will see how to save and load models in PyTorch. It is important because often we want to load the previously trained models to use:  \n",
    "1. in making predictions  \n",
    "2. to continue training  "
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this is to have the figures displayed in the\n",
    "# code and to set a format for the disaplying\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torch\n",
    "from torch import nn, 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",
    "# For the first time we have to set @download=True@ to download the dataset\n",
    "trainset = datasets.FashionMNIST('.\\\\F_MNIST\\\\', download=False, 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('.\\\\F_MNIST\\\\', download=False, train=False, transform=transform)\n",
    "testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=True)"
   ]
  },
  {
   "source": [
    "Here, we can see an image of the training set.  \n",
    "Each time we run the cell, we have that the subsequent batch of images is loaded.  \n",
    "And then, we are selecting the first image of the batch."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "<Figure size 432x288 with 1 Axes>",
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAc8AAAHPCAYAAAA1eFErAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuNCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8QVMy6AAAACXBIWXMAABYlAAAWJQFJUiTwAAALxUlEQVR4nO3dQW9U9xXG4Tv2jG1sS9iAgQoFiARRukmUKiUsuqqibqN+4DYqza5RI5IsAmlIMHEwYGIL24w9YzP9AlXF+z/UU4vn2R+dy4ztH3d1epPJpAMAXt/MtB8AAE4b8QSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEOq3Dv7xDx86xwLAqfbXL+72Wua8eQJASDwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJACHxBICQeAJASDwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJACHxBICQeAJASDwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJACHxBICQeAJASDwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJACHxBICQeAJASDwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACPWn/QDAfzY3N1eaH41Gb+hJOCnnVlebZ1dWVkq7H/z4Y2n+bePNEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJACHxBICQeAJASDwBICSeABBykowT0ev1mmcnk8kbfJKTNT8/3zz7u48+Ku3+51dfleYPDw9L86dV5Tv7zeXLU9t9ubjbSbKMN08ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJACHxBICQeAJASDwBIOSeJyfiNN/krLj6zjvNsxsbG6Xd/X7t1/vypUvNs3v7+6Xde3t7zbPVO6Tv3bzZPPvu9eul3d/du9c8+7c7d0q7yXjzBICQeAJASDwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACIknAIScJIP/4vz586X5yim29UePSrurBoNB8+yF4udWma+eJBuNRs2z33z7bWn31tZWaZ6T480TAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAi558n/vYWFhdL8ytmzzbPjo6PS7n/98ENpfpp2dnaaZw8ODkq7L66tNc/OzNTeCYbDYfPstG+wcnK8eQJASDwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBCTpLxWqpnnm7euNE8u7+/X9r95OnT5tnxeFza/baqniTb3t5unr1y5Upp997eXvNsr9cr7Z5MJqV5To43TwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEg5J7nCZudnW2e/fNnn5V2bz1/3jz74MGD0u579++X5nm77BZuag6Hw9LuxcXF5ln3ON8e3jwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAISfJQh9+8EFp/k+ffto8W712tLu72zy7/uhRafe51dXm2ReF5+66rjs6OirNc7o83twszS/MzzfPLi0tlXavrqxMbfc0nTlzpnn2542NN/gkr8+bJwCExBMAQuIJACHxBICQeAJASDwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASAkngAQOpX3PGdmas1/9epV8+za2lpp93B40Dx79+uvS7vvflObn5bbn3xSml9fX2+fLd4xpc3y8nLz7KWLF0u7K38fzp49W9rd77f/SR4Oh6Xds7OzzbPj8bi0ezAYNM8eT+lerzdPAAiJJwCExBMAQuIJACHxBICQeAJASDwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASB0Ku95Vu7tVb148aI0f//775tnn209K+1++fJl8+ztW7Wbmt/dv9c8+48vvyztvnnjRvPstatXS7sfFm6JVs3NzZXmz62uNs8uLS2VdlduS1ZVPrfqTc2K6mdWuclZ3T0ajZpnt3d2SrtbefMEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJACHxBICQeAJASDwBICSeABASTwAIiScAhKZ2kqzX6zXPVs/fHB0dNc9ubm6Wdm9tbTXPvnv9emn3s8LuX7d/Le2+fetW8+xfPv+8tPvx48fNs72Z2v8vl5eXm2fXLlwo7a6eJKv8nlR+v7uudnZwfn6+tLvy7x4MBqXdFcfHx6X5aZ56rHzm0+LNEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJACHxBICQeAJASDwBICSeABASTwAINd/zXFpaKi0+t7raPDs8OCjtnincGjw8PCztfrG72zxbvVNY+cw3fvmltLtyr++3779f2l35znYL31fXdd2lixebZw+KP+fj8bg0P1O4ZTqZTEq7K/dAR6NRaXflrmX1pma/335iuXpDtbL7zMKZ0u7tne3S/DR48wSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCEmm/QVM56dV3Xzc7ONs8eF85bdV3XjSpng4r/7rm5uebZp8+elXb3C5/50uJiaXflVNPz589LuyuntZaXl0u7K6fYqifoqhYL3/mgcN6q67ru5XDYPFs9jzU8aN89LDx3db56kqxyRq7yc951Xbe/v1+anwZvngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJACHxBICQeAJAqPno3mg8fpPPEbl27Vppvj/bfmuwX7xTePyq/a7laDQq7a48e/Ve36A/aJ6tfubjo/af1cod0q7runHh96RyX/FNzG9vbzfPVv7dXdd1CwsLzbO7u7ul3ZXv/Pcff1zaPTfXfsP18PCwtPtV4W/ThfMXSrvvfPH35tnNJ09Ku1t58wSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCEmm89Vc/f/PTwYftwZRbgf+Th+vq0H4ET4s0TAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJACHxBICQeAJASDwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJACHxBICQeAJASDwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJACHxBICQeAJASDwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJACHxBICQeAJASDwBICSeABASTwAIiScAhMQTAELiCQAh8QSAkHgCQEg8ASAkngAQEk8ACIknAITEEwBC4gkAIfEEgJB4AkCoN5lMpv0MAHCqePMEgJB4AkBIPAEgJJ4AEBJPAAiJJwCExBMAQuIJACHxBICQeAJASDwBICSeABASTwAI/Rsu1YeFQ/q4VgAAAABJRU5ErkJggg==\n"
     },
     "metadata": {
      "image/png": {
       "width": 231,
       "height": 231
      },
      "needs_background": "light"
     }
    }
   ],
   "source": [
    "image, label = next(iter(trainloader))\n",
    "# we take the first image of a batch of the train set\n",
    "helper.imshow(image[0,:]);"
   ]
  },
  {
   "source": [
    "## Train a Network\n",
    "The network defined in a previous notebook is reported in the module `fc_model`. Importing this, we can easily create a fully connected layer with `fc_model.Network` and train it with `fc_model.train`. We are going to save the model after it is trained. And then we will see how to load the model."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Epoch: 1/2..  Training Loss: 1.497..  Test Loss: 0.886..  Test Accuracy: 0.638\n",
      "Epoch: 1/2..  Training Loss: 0.994..  Test Loss: 0.757..  Test Accuracy: 0.692\n",
      "Epoch: 1/2..  Training Loss: 0.885..  Test Loss: 0.680..  Test Accuracy: 0.747\n",
      "Epoch: 1/2..  Training Loss: 0.812..  Test Loss: 0.641..  Test Accuracy: 0.752\n",
      "Epoch: 1/2..  Training Loss: 0.806..  Test Loss: 0.637..  Test Accuracy: 0.758\n",
      "Epoch: 1/2..  Training Loss: 0.793..  Test Loss: 0.664..  Test Accuracy: 0.750\n",
      "Epoch: 1/2..  Training Loss: 0.789..  Test Loss: 0.653..  Test Accuracy: 0.764\n",
      "Epoch: 1/2..  Training Loss: 0.795..  Test Loss: 0.633..  Test Accuracy: 0.770\n",
      "Epoch: 1/2..  Training Loss: 0.779..  Test Loss: 0.613..  Test Accuracy: 0.771\n",
      "Epoch: 1/2..  Training Loss: 0.780..  Test Loss: 0.618..  Test Accuracy: 0.759\n",
      "Epoch: 1/2..  Training Loss: 0.760..  Test Loss: 0.614..  Test Accuracy: 0.769\n",
      "Epoch: 1/2..  Training Loss: 0.752..  Test Loss: 0.617..  Test Accuracy: 0.777\n",
      "Epoch: 1/2..  Training Loss: 0.774..  Test Loss: 0.578..  Test Accuracy: 0.785\n",
      "Epoch: 1/2..  Training Loss: 0.746..  Test Loss: 0.580..  Test Accuracy: 0.784\n",
      "Epoch: 1/2..  Training Loss: 0.717..  Test Loss: 0.567..  Test Accuracy: 0.791\n",
      "Epoch: 1/2..  Training Loss: 0.751..  Test Loss: 0.565..  Test Accuracy: 0.799\n",
      "Epoch: 1/2..  Training Loss: 0.699..  Test Loss: 0.560..  Test Accuracy: 0.800\n",
      "Epoch: 1/2..  Training Loss: 0.763..  Test Loss: 0.551..  Test Accuracy: 0.802\n",
      "Epoch: 1/2..  Training Loss: 0.738..  Test Loss: 0.581..  Test Accuracy: 0.788\n",
      "Epoch: 1/2..  Training Loss: 0.721..  Test Loss: 0.562..  Test Accuracy: 0.793\n",
      "Epoch: 1/2..  Training Loss: 0.771..  Test Loss: 0.548..  Test Accuracy: 0.797\n",
      "Epoch: 1/2..  Training Loss: 0.712..  Test Loss: 0.548..  Test Accuracy: 0.797\n",
      "Epoch: 1/2..  Training Loss: 0.722..  Test Loss: 0.571..  Test Accuracy: 0.805\n",
      "Epoch: 2/2..  Training Loss: 0.715..  Test Loss: 0.549..  Test Accuracy: 0.811\n",
      "Epoch: 2/2..  Training Loss: 0.726..  Test Loss: 0.553..  Test Accuracy: 0.805\n",
      "Epoch: 2/2..  Training Loss: 0.698..  Test Loss: 0.556..  Test Accuracy: 0.801\n",
      "Epoch: 2/2..  Training Loss: 0.695..  Test Loss: 0.567..  Test Accuracy: 0.807\n",
      "Epoch: 2/2..  Training Loss: 0.738..  Test Loss: 0.540..  Test Accuracy: 0.805\n",
      "Epoch: 2/2..  Training Loss: 0.736..  Test Loss: 0.568..  Test Accuracy: 0.809\n",
      "Epoch: 2/2..  Training Loss: 0.738..  Test Loss: 0.555..  Test Accuracy: 0.794\n",
      "Epoch: 2/2..  Training Loss: 0.708..  Test Loss: 0.544..  Test Accuracy: 0.805\n",
      "Epoch: 2/2..  Training Loss: 0.651..  Test Loss: 0.546..  Test Accuracy: 0.810\n",
      "Epoch: 2/2..  Training Loss: 0.699..  Test Loss: 0.555..  Test Accuracy: 0.809\n",
      "Epoch: 2/2..  Training Loss: 0.724..  Test Loss: 0.538..  Test Accuracy: 0.822\n",
      "Epoch: 2/2..  Training Loss: 0.694..  Test Loss: 0.551..  Test Accuracy: 0.809\n",
      "Epoch: 2/2..  Training Loss: 0.705..  Test Loss: 0.543..  Test Accuracy: 0.802\n",
      "Epoch: 2/2..  Training Loss: 0.629..  Test Loss: 0.533..  Test Accuracy: 0.812\n",
      "Epoch: 2/2..  Training Loss: 0.728..  Test Loss: 0.521..  Test Accuracy: 0.822\n",
      "Epoch: 2/2..  Training Loss: 0.729..  Test Loss: 0.554..  Test Accuracy: 0.815\n",
      "Epoch: 2/2..  Training Loss: 0.698..  Test Loss: 0.522..  Test Accuracy: 0.824\n",
      "Epoch: 2/2..  Training Loss: 0.678..  Test Loss: 0.531..  Test Accuracy: 0.811\n",
      "Epoch: 2/2..  Training Loss: 0.710..  Test Loss: 0.539..  Test Accuracy: 0.814\n",
      "Epoch: 2/2..  Training Loss: 0.657..  Test Loss: 0.540..  Test Accuracy: 0.810\n",
      "Epoch: 2/2..  Training Loss: 0.695..  Test Loss: 0.525..  Test Accuracy: 0.819\n",
      "Epoch: 2/2..  Training Loss: 0.703..  Test Loss: 0.515..  Test Accuracy: 0.822\n",
      "Epoch: 2/2..  Training Loss: 0.690..  Test Loss: 0.526..  Test Accuracy: 0.814\n"
     ]
    }
   ],
   "source": [
    "# for how is defined the model, we have that the first two parameters of the functin @fc_model.Network@\n",
    "# takes the number of input nodes and the number of output nodes, then a list with the number of all the\n",
    "# layers in the hidden layers\n",
    "model = fc_model.Network(784, 10, [512, 256, 128])\n",
    "# the network ends with a logSoftMax layer, so we have to use a Negative Log Likelihood Loss function.\n",
    "criterion = nn.NLLLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.003)\n",
    "fc_model.train(model, trainloader, testloader, criterion, optimizer, epochs=2)"
   ]
  },
  {
   "source": [
    "## Saving and Loading Networks\n",
    "It is impractical to train a network every time we want to use it. Insted, we can save trained network and use them when we need.  \n",
    "So, we can use them later for predictions or for continuing to train.  \n",
    "The parameters for PyTorch networks are stored in a model's `state_dict()`.  \n",
    "We can see that the `state_dict()` contains the weights and biases matrices for each layer of our network."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "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\nThe 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())"
   ]
  },
  {
   "source": [
    "The simplest thing to do is simply save the state dict with `torch.save`. For example, we can save it in an object named `checkpoint.pth`."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), 'checkpoint.pth')"
   ]
  },
  {
   "source": [
    "Then, we can load the `state_dict()` with `torch.load()`."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "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())"
   ]
  },
  {
   "source": [
    "To load the state_dict to your network we do `model.load_state_dict(state_dict)`."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "metadata": {},
     "execution_count": 17
    }
   ],
   "source": [
    "model.load_state_dict(state_dict)"
   ]
  },
  {
   "source": [
    "Note: the loading of the previously saved model works only if we create a model which has the same architecture of the saved model.  \n",
    "Now, we will have a size mismatch error when trying to load the saved model to a model with a different architecture."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "output_type": "error",
     "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]).",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-18-d859c59ebec0>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      2\u001b[0m \u001b[0mmodel\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mfc_model\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mNetwork\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m784\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;36m10\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m[\u001b[0m\u001b[1;36m400\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;36m200\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;36m100\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      3\u001b[0m \u001b[1;31m# This will throw an error because the tensor sizes are wrong!\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 4\u001b[1;33m \u001b[0mmodel\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mload_state_dict\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mstate_dict\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;32m~\\ProjectsCode\\OtherStuff\\PyScripts\\UdacityDeepLearningPyTorch\\venv\\lib\\site-packages\\torch\\nn\\modules\\module.py\u001b[0m in \u001b[0;36mload_state_dict\u001b[1;34m(self, state_dict, strict)\u001b[0m\n\u001b[0;32m   1221\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1222\u001b[0m         \u001b[1;32mif\u001b[0m \u001b[0mlen\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0merror_msgs\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;33m>\u001b[0m \u001b[1;36m0\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m-> 1223\u001b[1;33m             raise RuntimeError('Error(s) in loading state_dict for {}:\\n\\t{}'.format(\n\u001b[0m\u001b[0;32m   1224\u001b[0m                                self.__class__.__name__, \"\\n\\t\".join(error_msgs)))\n\u001b[0;32m   1225\u001b[0m         \u001b[1;32mreturn\u001b[0m \u001b[0m_IncompatibleKeys\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mmissing_keys\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0munexpected_keys\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;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": [
    "# a model with a different architecture\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)"
   ]
  },
  {
   "source": [
    "This means that we have to build the model exactly as it was when trained. Information about the model architecture has to be saved in the checkpoint, along with the state dict.  \n",
    "To do this, we can build a dictionary with all the information about the architecture of the network we are going to save (even though, you can use another type of object).  "
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "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')"
   ]
  },
  {
   "source": [
    "Now, the checkpoint model has all the information to rebuild the train model.  \n",
    "We create a custom function to load the checkpoint`load_checkpoint`."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "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": 22,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "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": []
  }
 ]
}