{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "0386b88759cd8dda9c6211fc1ef856a5",
     "grade": false,
     "grade_id": "cell-4b3a34920f4c7395",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "# Part 1: Training and evaluating simple CNNs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "a06e331b63eb9f8b5511ee8e90ae3dc6",
     "grade": false,
     "grade_id": "cell-2852eef932feefd7",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "__Before starting, we recommend you enable GPU acceleration if you're running on Colab.__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "84a38ee23b261ec917497a78bd147d1f",
     "grade": false,
     "grade_id": "cell-6f99954a67e07cc3",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "# Execute this code block to install dependencies when running on colab\n",
    "try:\n",
    "    import torch\n",
    "except:\n",
    "    from os.path import exists\n",
    "    from wheel.pep425tags import get_abbr_impl, get_impl_ver, get_abi_tag\n",
    "    platform = '{}{}-{}'.format(get_abbr_impl(), get_impl_ver(), get_abi_tag())\n",
    "    cuda_output = !ldconfig -p|grep cudart.so|sed -e 's/.*\\.\\([0-9]*\\)\\.\\([0-9]*\\)$/cu\\1\\2/'\n",
    "    accelerator = cuda_output[0] if exists('/dev/nvidia0') else 'cpu'\n",
    "\n",
    "    !pip install -q http://download.pytorch.org/whl/{accelerator}/torch-1.0.0-{platform}-linux_x86_64.whl torchvision\n",
    "\n",
    "try: \n",
    "    import torchbearer\n",
    "except:\n",
    "    !pip install torchbearer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "ce847410c1ef0a28d490008cedf86b4a",
     "grade": false,
     "grade_id": "cell-d0369c9de7c5a89e",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "## Simple Convolutional Neural Network for MNIST\n",
    "\n",
    "Now that we have seen how to load the MNIST dataset and train a simple multi-layer perceptron model on it, we can now start to develop a more sophisticated convolutional neural network or CNN model. PyTorch provides a lot of capability for creating CNNs, and includes a large number of layer types and activation functions. In this part of the lab we will create a simple CNN for MNIST that demonstrates how to use all of the aspects of a modern CNN implementation, including Convolutional layers, Pooling layers and Dropout layers.\n",
    "\n",
    "The first step is to import the classes and functions needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "175dd97d30f646eb74ad9e21427a2c0d",
     "grade": false,
     "grade_id": "cell-b7a634722fea19b2",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "# automatically reload external modules if they change\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torchvision.transforms as transforms\n",
    "import torchbearer\n",
    "from torch import nn\n",
    "from torch import optim\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision.datasets import MNIST\n",
    "from torchbearer import Trial"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "dbd5411f763365d82fece69745457168",
     "grade": false,
     "grade_id": "cell-36d240a1fee1f6f2",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "Again, initialise the random number generator to a constant seed value for reproducibility of results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "b422e04bbefbf78800d34b845c1397d5",
     "grade": false,
     "grade_id": "cell-4d617d7c9adb1794",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "# fix random seed for reproducibility\n",
    "seed = 7\n",
    "torch.manual_seed(seed)\n",
    "torch.backends.cudnn.deterministic = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "1caef4794d2c1a7d0a0ca69b5c8c3669",
     "grade": false,
     "grade_id": "cell-f5f854670ab50aed",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "Next we need to load the MNIST dataset. Unlike with the MLP example we looked at previously, we don't want to flatten the images into vectors, although we do still want to convert the PIL images to tensors.\n",
    "\n",
    "In PyTorch, images are represented as tensors with dimensions `[pixels][height][width]`. In the case of RGB, the first dimension, pixels, would be 3 for the red, green and blue components. In the case of MNIST where the pixel values are greyscale, the pixel dimension is set to 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "29254ea7761e47b036e1bf7e6e6f8f00",
     "grade": false,
     "grade_id": "cell-fe9cacfcc16dcddf",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "# convert each image to tensor format\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor()  # convert to tensor\n",
    "])\n",
    "\n",
    "# load data\n",
    "trainset = MNIST(\".\", train=True, download=True, transform=transform)\n",
    "testset = MNIST(\".\", train=False, download=True, transform=transform)\n",
    "\n",
    "# create data loaders\n",
    "trainloader = DataLoader(trainset, batch_size=128, shuffle=True)\n",
    "testloader = DataLoader(testset, batch_size=128, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "aa79db913111d98503bf94cc0f4c77a8",
     "grade": false,
     "grade_id": "cell-f094c93f38605686",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "Next we define our neural network model.\n",
    "\n",
    "Convolutional neural networks are more complex than standard multi-layer perceptrons, so we will start by using a simple structure to begin with that uses all of the elements required for near state of the art results. The network architecture is summarised below:\n",
    "\n",
    "1. The first hidden layer is a convolutional layer called a `Convolution2D`. The layer has 32 feature maps, which with the size of 5×5 and a rectified linear unit activation function.\n",
    "2. Next we define a pooling layer that takes the max called `MaxPooling2D`. It is configured with a pool size of 2×2.\n",
    "3. The next layer is a regularization layer using dropout called `Dropout`. It is configured to randomly exclude 20% of neurons in the layer in order to reduce overfitting.\n",
    "4. Next is an operation that flattens or reshapes the tensor to a vector. It allows the output to be processed by standard fully connected layers.\n",
    "5. Next a fully connected layer with 128 neurons and rectifier linear unit activation function.\n",
    "6. Finally, the output layer has 10 neurons for the 10 classes.\n",
    "\n",
    "Just like with the MLP we create a PyTorch `nn.Module` subclass to define our network model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "d978becfd60240a69ccfcb346d9fabdb",
     "grade": false,
     "grade_id": "cell-e36a90ba5353bd77",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "class SimpleCNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(SimpleCNN, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 32, (5, 5), padding=0)\n",
    "        self.fc1 = nn.Linear(32 * 12**2, 128)\n",
    "        self.fc2 = nn.Linear(128, 10)\n",
    "            \n",
    "    def forward(self, x):\n",
    "        out = self.conv1(x)\n",
    "        out = F.relu(out)\n",
    "        out = F.max_pool2d(out, (2,2))\n",
    "        out = F.dropout(out, 0.2)\n",
    "        out = out.view(out.shape[0], -1)\n",
    "        out = self.fc1(out)\n",
    "        out = F.relu(out)\n",
    "        out = self.fc2(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "80b333944df484312ce0c3804eca5e6b",
     "grade": false,
     "grade_id": "cell-7483d476ef9f0de6",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "Note in the above code when we use the `view` method to _flatten_ the tensor we must preserve the batch dimension of the data; passing `-1` as the second argument tells view to compute the size of the dimension represented by the `-1` automatically. The input to the forward method has shape `[batch_size][pixels][height][width]` and we want the final output to be `[batch_size][num_classes=10]`. \n",
    "\n",
    "__Answer the following questions (enter the answer in the box below each one):__\n",
    "\n",
    "__1.__ What is the shape of the tensor output by the `conv1` layer?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "checksum": "a8e7f96c4270925d62470b6875acaa0d",
     "grade": true,
     "grade_id": "cell-85c7d9e142f0f5a8",
     "locked": false,
     "points": 1,
     "schema_version": 1,
     "solution": true
    }
   },
   "source": [
    "YOUR ANSWER HERE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "9ba2f36afc0dd997780d17b4d9ce5e04",
     "grade": false,
     "grade_id": "cell-423f347f09d8ecd9",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "__2.__ If the `conv1` layer had a kernel shape of `(7,3)` what would the output shape be?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "checksum": "5c8923855a9c4f95ffbd658bfaf8201a",
     "grade": true,
     "grade_id": "cell-5bded1643d1a5530",
     "locked": false,
     "points": 2,
     "schema_version": 1,
     "solution": true
    }
   },
   "source": [
    "YOUR ANSWER HERE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "e690b390d913ddf3d089a55a8c9b9887",
     "grade": false,
     "grade_id": "cell-369d69d982d860ef",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "As before, the model is trained using cross-entropy loss and the ADAM gradient descent algorithm. The CNN is fit over 10 epochs with a batch size of 128. We use torchbearer to minimise the training and evaluation code we need to write:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "666b5f1daf8391d5c1a1ab426caddf90",
     "grade": false,
     "grade_id": "cell-03a034a2e45728dd",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "# build the model\n",
    "model = SimpleCNN()\n",
    "\n",
    "# define the loss function and the optimiser\n",
    "loss_function = nn.CrossEntropyLoss()\n",
    "optimiser = optim.Adam(model.parameters())\n",
    "\n",
    "device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n",
    "trial = Trial(model, optimiser, loss_function, metrics=['loss', 'accuracy']).to(device)\n",
    "trial.with_generators(trainloader, test_generator=testloader)\n",
    "trial.run(epochs=10)\n",
    "results = trial.evaluate(data_key=torchbearer.TEST_DATA)\n",
    "print(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "671150370dacaaf1a06bf5c07f288cbf",
     "grade": false,
     "grade_id": "cell-b839b934e04f89c0",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "You should see that the network achieves an accuracy of 98.91% which is better than the simple multi-layer perceptron model we tried previously.\n",
    "\n",
    "## An improved CNN\n",
    "\n",
    "Now that we have seen how to create a simple CNN, let’s take a look at a model capable of close to state of the art results. This time you will implement a large CNN architecture with additional convolutional, max pooling layers and fully connected layers. The network topology of the model is summarised as follows:\n",
    "\n",
    "1. Convolutional layer with 30 feature maps of size 5×5 and ReLU activation.\n",
    "2. Pooling layer taking the max over 2×2 patches.\n",
    "3. Convolutional layer with 15 feature maps of size 3×3 and ReLU activation.\n",
    "4. Pooling layer taking the max over 2×2 patches.\n",
    "5. Dropout layer with a probability of 20%.\n",
    "6. Flatten layer.\n",
    "7. Fully connected layer with 128 neurons and ReLU activation.\n",
    "8. Fully connected layer with 50 neurons and ReLU activation.\n",
    "9. Linear output layer.\n",
    "\n",
    "__Complete the implementation of the BetterCNN module which implements the above network in the following code block.__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "checksum": "3a544d1f48cb9b1af278acbebdf391be",
     "grade": false,
     "grade_id": "cell-110a1189124033d9",
     "locked": false,
     "schema_version": 1,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "import torch \n",
    "import torch.nn.functional as F\n",
    "from torch import nn\n",
    "\n",
    "# Model Definition\n",
    "class BetterCNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(BetterCNN, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 30, (5, 5), padding=0)\n",
    "        self.conv2 = nn.Conv2d(30, 15, (3, 3), padding=0)\n",
    "        self.fc1 = nn.Linear(15 * 5**2, 128)\n",
    "        self.fc2 = nn.Linear(128, 50)\n",
    "        self.fc3 = nn.Linear(50, 10)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        # YOUR CODE HERE\n",
    "        raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "9fa81f5c9c2759a8b1fbc372028497c3",
     "grade": false,
     "grade_id": "cell-ad9d70fbb2d8fc5f",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "Once you've completed the implementation, save the file and run the following code block to train and evaluate the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "5cf30cdbb30e88f01a91f8eebfc8246c",
     "grade": true,
     "grade_id": "cell-da90b81391df00b2",
     "locked": true,
     "points": 5,
     "schema_version": 1,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "#reset the data loaders\n",
    "torch.manual_seed(seed)\n",
    "trainloader = DataLoader(trainset, batch_size=128, shuffle=True)\n",
    "testloader = DataLoader(testset, batch_size=128, shuffle=True)\n",
    "\n",
    "# build the model\n",
    "model = BetterCNN()\n",
    "\n",
    "# define the loss function and the optimiser\n",
    "loss_function = nn.CrossEntropyLoss()\n",
    "optimiser = optim.Adam(model.parameters())\n",
    "\n",
    "device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n",
    "trial = Trial(model, optimiser, loss_function, metrics=['loss', 'accuracy']).to(device)\n",
    "trial.with_generators(trainloader, test_generator=testloader)\n",
    "trial.run(epochs=10)\n",
    "results = trial.evaluate(data_key=torchbearer.TEST_DATA)\n",
    "print(results)\n",
    "\n",
    "assert results['test_acc'] > 0.99\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "c2baca3abcdbfc058a657e165a1cf8e3",
     "grade": false,
     "grade_id": "cell-1c2475b8550316a4",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "If correctly implemented you should see this slightly larger model achieves the respectable classification accuracy of 99.13%.\n",
    "\n",
    "## Saving models\n",
    "\n",
    "Being able to train a model is fine, but in practice once we've trained the model we probably want to save the result so we can reuse it at a later time. PyTorch makes saving the model easy using the `torch.save(state, filepath)` function. This will save the weights of the model so they can be loaded into a new instance at a later point. \n",
    "\n",
    "__Run the following code to save the weights for use in the next part of the lab.__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "4895367de85f19ae1bb89e90ad300e70",
     "grade": false,
     "grade_id": "cell-6d81f64d4d65498c",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "#save the trained model weights\n",
    "torch.save(model.state_dict(), \"./bettercnn.weights\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "5152e5288466c2af37eb11d3393c5446",
     "grade": false,
     "grade_id": "cell-35c249867f1ec1d2",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "__If you are running on Colab, run the following to download the weights to the local machine:__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "c75f37cdabf24f935575dcb026f3c6f3",
     "grade": false,
     "grade_id": "cell-957d645048f13f7a",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "from google.colab import files\n",
    "files.download('bettercnn.weights')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}