mnist/2_mnist_model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PyTorch MNIST example dissected\n",
    "\n",
    "In this notebook we'll explore the components of the\n",
    "[PyTorch MNIST example](https://github.com/pytorch/examples/tree/master/mnist)\n",
    "one-by-one.\n",
    "\n",
    "* Part 1: [Loading the data](1_mnist_load.ipynb)\n",
    "* Part 2: [Model components and forward propagation](2_mnist_model.ipynb##2-Building-the-model) <-- **you are here**\n",
    "* Part 3: [Autodiff and backpropagation](3_mnist_backprop.ipynb)\n",
    "* Part 4: [Training the model](4_mnist_train.ipynb)\n",
    "* Part 5: [Visualizing the results](5_mnist_visualize.ipynb)\n",
    "\n",
    "In this section:\n",
    "\n",
    "* [2.0 Input Layer](#2.0-Input-Layer)\n",
    "* [2.1 Conv2d](#2.1-Conv2d)\n",
    "  * [2.1.1 Sliding window](#2.1.1-Sliding-window)\n",
    "  * [2.1.2 Aside: Gradient computation](#2.1.2-Aside:-Gradient-computation)\n",
    "* [2.2 Max pooling](#2.2-Max-pooling)\n",
    "* [2.3 ReLU](#2.3-ReLU)\n",
    "* [2.4 Layer 2](#2.4-Layer-2)\n",
    "* [2.5 2d Dropout](#2.5-2d-Dropout)\n",
    "  * [2.5.1 Aside: Dropout techniques](#2.5.1-Aside:-Dropout-techniques)\n",
    "* [2.6 Layers 3 and 4](#2.6-Layers-3-and-4)\n",
    "  * [2.6.1 Aside: Dimensions](#2.6.1-Aside:-Dimensions)\n",
    "  * [2.6.2 Linear unit](#2.6.2-Linear-unit)\n",
    "* [2.7 Softmax](#2.7-Softmax)\n",
    "* [2.8 Conclusion](#2.8-Conclusion)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2 Building the model\n",
    "\n",
    "In this notebook we'll explore the components of `torch.nn`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll start with copying the model from the original [PyTorch MNIST example](https://github.com/pytorch/examples/tree/master/mnist) verbatim:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)\n",
    "        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n",
    "        self.conv2_drop = nn.Dropout2d()\n",
    "        self.fc1 = nn.Linear(320, 50)\n",
    "        self.fc2 = nn.Linear(50, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(F.max_pool2d(self.conv1(x), 2))\n",
    "        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n",
    "        x = x.view(-1, 320)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.dropout(x, training=self.training)\n",
    "        x = self.fc2(x)\n",
    "        return F.log_softmax(x, dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That is, any class derived from `nn.Module` should implement the `.forward()` method; everything else is taken care of automagically by the auto-differentiation framework. Pretty neat.\n",
    "\n",
    "The original `x` argument of the `.forward()` method is the **batch** of features generated by the `DataLoader` (see [Part 1](1_mnist_load.ipynb#1.4-Dataset-and-DataLoader) for details). That is, for the MNIST dataset and batch size of 64, the dimensions of the input `x` are\n",
    "\n",
    "    torch.Size([64, 1, 28, 28])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.0 Input Layer\n",
    "\n",
    "First layer of our neural net is defined as\n",
    "\n",
    "```python\n",
    "x = F.relu(F.max_pool2d(self.conv1(x), 2))\n",
    "```\n",
    "\n",
    "The combination of 2D convolution, max pooling, and ReLU activation function if very typical for learning image models. Below we will look at each of these three components in detail."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Conv2d\n",
    "\n",
    "First, we feed the input batch through 2D convolution:\n",
    "\n",
    "```python\n",
    "self.conv1(x)\n",
    "```\n",
    "\n",
    " this convolution is defined with 5x5 kernel and `10` output layers:\n",
    "\n",
    "```python\n",
    "self.conv1 = nn.Conv2d(1, 10, kernel_size=5)\n",
    "```\n",
    "\n",
    "Recall that our input has 1 layer, because it is a greyscale image; hence `1` in `Conv2d()` constructor parameters.\n",
    "\n",
    "Below we will create an instance of `Conv2d` and see what it actually does to the data. We start with a tiny 3x3 input tensor and 3x3 convolution with 4 output layers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[0., 1., 1.],\n",
       "          [1., 0., 1.],\n",
       "          [1., 1., 0.]]]])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = 1 - torch.eye(3, dtype=torch.float32).reshape((1, 1, 3, 3))\n",
    "data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that NN layers operate on *batches*, hence input data dimensions of [batch of size 1 x 1 color layer x height 3 x width 3]: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 1, 3, 3])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Conv2d(1, 4, kernel_size=(3, 3), stride=(1, 1))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conv = nn.Conv2d(1, 4, kernel_size=3)\n",
    "conv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[-0.2512]],\n",
       "\n",
       "         [[-0.2286]],\n",
       "\n",
       "         [[-0.2745]],\n",
       "\n",
       "         [[ 0.1661]]]], grad_fn=<ThnnConv2DBackward>)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conv_result = conv(data)\n",
    "conv_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 4, 1, 1])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conv_result.size()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So our convolution converts a 1x3x3 patch of (greyscale) image into a 4x1x1 column of weird numbers. Where do those numbers come from?\n",
    "\n",
    "Let's take a look at the `Conv2d` internals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "(params, bias) = list(conv.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[[[ 0.0234,  0.0275, -0.2540],\n",
       "          [ 0.1430,  0.2409,  0.0962],\n",
       "          [ 0.2230, -0.2230,  0.2775]]],\n",
       "\n",
       "\n",
       "        [[[-0.3157, -0.3143, -0.1466],\n",
       "          [ 0.1508, -0.0177, -0.1226],\n",
       "          [-0.0785,  0.0568, -0.1158]]],\n",
       "\n",
       "\n",
       "        [[[-0.0449, -0.0043, -0.1664],\n",
       "          [ 0.2196, -0.2802, -0.0740],\n",
       "          [ 0.2470, -0.2914,  0.3084]]],\n",
       "\n",
       "\n",
       "        [[[ 0.2295, -0.2852,  0.2646],\n",
       "          [ 0.2382,  0.1396, -0.2646],\n",
       "          [ 0.1704,  0.1809, -0.2728]]]], requires_grad=True)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([-0.2638,  0.2257, -0.2050, -0.1382], requires_grad=True)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bias"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These are randomly initialized parameters of our convolution. The second one is a bias vector. (It is possible to have a convolution *without* a bias - that's `bias=False` parameter in the constructor).\n",
    "\n",
    "So what is exactly the computation that `Conv2d` performs on the data during feed forward?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.2512, -0.2286, -0.2745,  0.1661], grad_fn=<ThAddBackward>)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(params * data).sum(dim=(1, 2, 3)) + bias"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Woo hoo, we get the same results!\n",
    "\n",
    "That is, for each of 4 output layers of the convolution there is a 3x3 matrix of parameters that we multiply by a 3x3 patch of the input image *element-wise*, and then sum up all elements of that product.\n",
    "\n",
    "Let's do it again, just for one layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.0234,  0.0275, -0.2540],\n",
       "         [ 0.1430,  0.2409,  0.0962],\n",
       "         [ 0.2230, -0.2230,  0.2775]]], grad_fn=<SelectBackward>)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "params[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[ 0.0000,  0.0275, -0.2540],\n",
       "          [ 0.1430,  0.0000,  0.0962],\n",
       "          [ 0.2230, -0.2230,  0.0000]]]], grad_fn=<ThMulBackward>)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "params[0] * data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(-0.2512, grad_fn=<ThAddBackward>)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(params[0] * data).sum() + bias[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each of 4 layers of the convolution parameters can be seen as a 3x3 patch of an image that we learn to recognize in the input. Right now it is random, but as we learn (i.e. backpropagate the gradient through it), it will morph into something like, say, a vertical line, or a diagonal gradient. Having 4 layers means that our convolution learns to recognize 4 different image patterns. Our original MNIST model has 10 such layers in its first convolution.\n",
    "\n",
    "Having the bias parameter for each layer allows us to learn image patterns of different intensity."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.1.1 Sliding window\n",
    "\n",
    "In our toy example above input image is the same size as the convolution (3x3). Let's apply our convolution to a larger image now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[0., 1., 1., 1., 1.],\n",
       "          [1., 0., 1., 1., 1.],\n",
       "          [1., 1., 0., 1., 1.],\n",
       "          [1., 1., 1., 0., 1.],\n",
       "          [1., 1., 1., 1., 0.]]]])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data2 = 1 - torch.eye(5, dtype=torch.float32).reshape((1, 1, 5, 5))\n",
    "data2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[-0.2512,  0.3706,  0.0675],\n",
       "          [ 0.1668, -0.2512,  0.3706],\n",
       "          [ 0.5445,  0.1668, -0.2512]],\n",
       "\n",
       "         [[-0.2286, -0.8856, -0.5994],\n",
       "          [-0.2411, -0.2286, -0.8856],\n",
       "          [-0.5313, -0.2411, -0.2286]],\n",
       "\n",
       "         [[-0.2745, -0.2194, -0.5382],\n",
       "          [-0.2129, -0.2745, -0.2194],\n",
       "          [-0.1249, -0.2129, -0.2745]],\n",
       "\n",
       "         [[ 0.1661, -0.1567,  0.0920],\n",
       "          [ 0.8123,  0.1661, -0.1567],\n",
       "          [-0.0021,  0.8123,  0.1661]]]], grad_fn=<ThnnConv2DBackward>)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conv_result2 = conv(data2)\n",
    "conv_result2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the elements at `[0,0]` in each of 4 layers are exactly the same as we had from our 3x3 image:  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.2512, -0.2286, -0.2745,  0.1661]], grad_fn=<SelectBackward>)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conv_result2[:,:,0,0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[-0.2512]],\n",
       "\n",
       "         [[-0.2286]],\n",
       "\n",
       "         [[-0.2745]],\n",
       "\n",
       "         [[ 0.1661]]]], grad_fn=<ThnnConv2DBackward>)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conv(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So the 2d convolution applies the 3x3 sliding window with the same parameters to each 3x3 patch of the input. In other words, on the forward pass it converts a 1x3x3 patch of input into 4x1x1 column output, and builds a new 4-layer tensor out of these elements.\n",
    "\n",
    "By default, 2d convolution has a **stride** of `1` and **padding** of `0`. That means, we'll move our sliding window by 1 pixel, starting from the [0,0] element and never going *outside* of the input - i.e. applying it first to"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[0., 1., 1.],\n",
       "          [1., 0., 1.],\n",
       "          [1., 1., 0.]]]])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data2[:, :, 0:3, 0:3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "then to"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[1., 1., 1.],\n",
       "          [0., 1., 1.],\n",
       "          [1., 0., 1.]]]])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data2[:, :, 0:3, 1:4]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and so on, all the way to `data2[:, :, 2:5, 2:5]`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.1.2 Aside: Gradient computation\n",
    "\n",
    "Note that tensors holding the convolution parameters have property `requires_grad=True`.\n",
    "\n",
    "Likewise, the results of the convolution are tensors with `grad_fn` property set.\n",
    "\n",
    "This is where the auto-differentiation magic of PyTorch happens. Later in this course we will dedicate the whole section to discuss it (see [Backpropagation](3_mnist_backprop.ipynb)). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Max pooling\n",
    "\n",
    "Next, we apply max pooling to the convolution results, i.e.\n",
    "\n",
    "```python\n",
    "F.max_pool2d(self.conv1(x), 2)\n",
    "```\n",
    "\n",
    "Here it runs 2x2 sliding window through the data and picks the largest values from each layer within that window. Note that unlike the convolution, by default max pooling windows do *not* overlap.\n",
    "\n",
    "Let's see how it works on some toy data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[11., 12., 13., 14., 15., 16.],\n",
       "         [21., 22., 23., 24., 25., 26.],\n",
       "         [31., 32., 33., 34., 35., 36.],\n",
       "         [41., 42., 43., 44., 45., 46.],\n",
       "         [51., 52., 53., 54., 55., 56.],\n",
       "         [61., 62., 63., 64., 65., 66.]]])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.linspace(10,60,6).reshape((1,6,1)).repeat(1,1,6) + torch.linspace(1,6,6).reshape((1,1,6))\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[22., 24., 26.],\n",
       "         [42., 44., 46.],\n",
       "         [62., 64., 66.]]])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "F.max_pool2d(x, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[33., 36.],\n",
       "         [63., 66.]]])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "F.max_pool2d(x, 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That is, max pooling effectively scales down the input proportionally to the kernel size. \n",
    "\n",
    "If input dimension is not a multiple of max pooling kernel size, we may lose some information at the edges, e.g."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[55.]]])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "F.max_pool2d(x, 5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's apply it to the multilayer tensor that is the result of our 2d convolution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[-0.2512,  0.3706,  0.0675],\n",
       "          [ 0.1668, -0.2512,  0.3706],\n",
       "          [ 0.5445,  0.1668, -0.2512]],\n",
       "\n",
       "         [[-0.2286, -0.8856, -0.5994],\n",
       "          [-0.2411, -0.2286, -0.8856],\n",
       "          [-0.5313, -0.2411, -0.2286]],\n",
       "\n",
       "         [[-0.2745, -0.2194, -0.5382],\n",
       "          [-0.2129, -0.2745, -0.2194],\n",
       "          [-0.1249, -0.2129, -0.2745]],\n",
       "\n",
       "         [[ 0.1661, -0.1567,  0.0920],\n",
       "          [ 0.8123,  0.1661, -0.1567],\n",
       "          [-0.0021,  0.8123,  0.1661]]]], grad_fn=<ThnnConv2DBackward>)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conv_result2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[ 0.5445]],\n",
       "\n",
       "         [[-0.2286]],\n",
       "\n",
       "         [[-0.1249]],\n",
       "\n",
       "         [[ 0.8123]]]], grad_fn=<MaxPool2DWithIndicesBackward>)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "F.max_pool2d(conv_result2, 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Recall that each of 4 layers correspond to some image pattern represented by the convolution parameters. Max pooling, then, allows us to learn the location of each pattern in the input image."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 ReLU\n",
    "\n",
    "Finally, we apply the Rectified Linear Unit activation function to the output:\n",
    "\n",
    "```python\n",
    "x = F.relu(F.max_pool2d(self.conv1(x), 2))\n",
    "```\n",
    "\n",
    "and that concludes the first layer of our neural network.\n",
    "\n",
    "`relu()` function simply trims the negative elements of the tensor to zero, e.g."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0, 0, 0, 1, 2])"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "F.relu(torch.tensor([-2, -1, 0, 1, 2]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This way, we will not propagate the gradient back to the layers that do not have positive correlation with certain patches of the image."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.4 Layer 2\n",
    "\n",
    "Next layer of our framework is very similar to the first one:\n",
    "\n",
    "```python\n",
    "x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n",
    "```\n",
    "   \n",
    "Where `conv2` is defined as\n",
    "\n",
    "```python\n",
    "self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n",
    "```\n",
    "\n",
    "So we have the same sequence of `Conv2d`, `max_pool2d`, and `relu`.\n",
    "\n",
    "To match the output of the first layer, `conv2` has `10` input dimensions, and `20` output ones. That means, on this stage out neural net will learn 20 different 5x5 patterns, each being a mosaic of 5x5 image components from the layer below. We'll try to visualize these parameters later, after we train our network."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.5 2d Dropout\n",
    "\n",
    "One new component that we encounter in the second layer is **2d dropout**, `conv2_drop`. It is defined as\n",
    "\n",
    "```python\n",
    "self.conv2_drop = nn.Dropout2d()\n",
    "```\n",
    "\n",
    "Let us see what it does."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dropout2d(p=0.3333333333333333)"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "drop = nn.Dropout2d(p=1./3.)\n",
    "drop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Parameter `p` specifies the probability of setting random *layers* of the input tensor to zero. The [Deep Learning book](http://www.deeplearningbook.org) suggests `p=0.8` for input layers and `p=0.5` for hidden layers of the net.\n",
    "\n",
    "For each new *batch*, `Dropout2d` will pick different layers at random."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[-0.2512,  0.3706,  0.0675],\n",
       "          [ 0.1668, -0.2512,  0.3706],\n",
       "          [ 0.5445,  0.1668, -0.2512]],\n",
       "\n",
       "         [[-0.2286, -0.8856, -0.5994],\n",
       "          [-0.2411, -0.2286, -0.8856],\n",
       "          [-0.5313, -0.2411, -0.2286]],\n",
       "\n",
       "         [[-0.2745, -0.2194, -0.5382],\n",
       "          [-0.2129, -0.2745, -0.2194],\n",
       "          [-0.1249, -0.2129, -0.2745]],\n",
       "\n",
       "         [[ 0.1661, -0.1567,  0.0920],\n",
       "          [ 0.8123,  0.1661, -0.1567],\n",
       "          [-0.0021,  0.8123,  0.1661]]]], grad_fn=<ThnnConv2DBackward>)"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conv_result2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[[-0.3768,  0.5558,  0.1012],\n",
      "          [ 0.2502, -0.3768,  0.5558],\n",
      "          [ 0.8168,  0.2502, -0.3768]],\n",
      "\n",
      "         [[-0.3430, -1.3284, -0.8991],\n",
      "          [-0.3616, -0.3430, -1.3284],\n",
      "          [-0.7970, -0.3616, -0.3430]],\n",
      "\n",
      "         [[-0.4118, -0.3291, -0.8073],\n",
      "          [-0.3193, -0.4118, -0.3291],\n",
      "          [-0.1873, -0.3193, -0.4118]],\n",
      "\n",
      "         [[ 0.0000, -0.0000,  0.0000],\n",
      "          [ 0.0000,  0.0000, -0.0000],\n",
      "          [-0.0000,  0.0000,  0.0000]]]], grad_fn=<FeatureDropoutBackward>)\n"
     ]
    }
   ],
   "source": [
    "drop_result = drop(conv_result2)\n",
    "\n",
    "print(drop_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that `Dropout2d` not only zeroes some (random) layers, but also rescales the remaining non-zero values by `1/(1-p)`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[1.5000, 1.5000, 1.5000],\n",
       "          [1.5000, 1.5000, 1.5000],\n",
       "          [1.5000, 1.5000, 1.5000]],\n",
       "\n",
       "         [[1.5000, 1.5000, 1.5000],\n",
       "          [1.5000, 1.5000, 1.5000],\n",
       "          [1.5000, 1.5000, 1.5000]],\n",
       "\n",
       "         [[1.5000, 1.5000, 1.5000],\n",
       "          [1.5000, 1.5000, 1.5000],\n",
       "          [1.5000, 1.5000, 1.5000]],\n",
       "\n",
       "         [[0.0000, 0.0000, 0.0000],\n",
       "          [0.0000, 0.0000, 0.0000],\n",
       "          [0.0000, 0.0000, 0.0000]]]], grad_fn=<DivBackward1>)"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "drop_result / conv_result2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.5.1 Aside: Dropout techniques\n",
    "\n",
    "Dropout is a popular regularization technique that promotes independence between elements of the neural net. Since the dropout mask changes from one batch to the next, we effectively create an *ensemble* of models that share some of their parameters.\n",
    "\n",
    "Unlike `Dropout2d`, regular `Dropout` class zeroes out neurons at random, e.g."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[-0.0000,  0.0000,  0.1349],\n",
       "          [ 0.0000, -0.5024,  0.0000],\n",
       "          [ 0.0000,  0.3337, -0.5024]],\n",
       "\n",
       "         [[-0.0000, -0.0000, -1.1988],\n",
       "          [-0.0000, -0.4573, -1.7712],\n",
       "          [-1.0627, -0.0000, -0.4573]],\n",
       "\n",
       "         [[-0.5490, -0.4388, -1.0764],\n",
       "          [-0.0000, -0.5490, -0.0000],\n",
       "          [-0.2497, -0.0000, -0.0000]],\n",
       "\n",
       "         [[ 0.0000, -0.0000,  0.1840],\n",
       "          [ 0.0000,  0.0000, -0.0000],\n",
       "          [-0.0043,  0.0000,  0.3321]]]], grad_fn=<DropoutBackward>)"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nn.Dropout(p=0.5)(conv_result2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This does not work well for convolutional nets, where each pixel of the input is strongly correlated with its immediate neigbors. For images, zeroing out random neurons will just slow down the learning process. Layer-wise dropout implemented in `Dropout2d` is much more efficient in such case."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Besides the `Dropout*` classes, PyTorch also has dropout functions, e.g."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0., 0., 4., 4., 4.])"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "F.dropout(torch.ones(5), p=0.75, training=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the `training=True` flag. We need it because the dropout mask is applied for the training phase only, and *not* for testing or production. When it is set to `False`, which is the default, `dropout()` will pass the data through unchanged. `nn.Dropout*` classes work the same way too, but we don't have to pass the `training` flag explicitly - it is already there as part of `nn.Module` base class."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.6 Layers 3 and 4\n",
    "\n",
    "Layers that follow the two convolutional layers are defined as:\n",
    "\n",
    "```python\n",
    "self.fc1 = nn.Linear(320, 50)\n",
    "self.fc2 = nn.Linear(50, 10)\n",
    "\n",
    "# ...\n",
    "\n",
    "x = x.view(-1, 320)\n",
    "x = F.relu(self.fc1(x))\n",
    "x = F.dropout(x, training=self.training)\n",
    "x = self.fc2(x)\n",
    "```\n",
    "\n",
    "i.e., first we have a simple linear layer from 320 to 50 elements followed by a familiar ReLU activation function;\n",
    "after it goes a 50% dropout, and another linear layer with 10 output features (one for each MNIST digit).\n",
    "\n",
    "#### 2.6.1 Aside: Dimensions\n",
    "\n",
    "Where's that dimension of `320` come from? Let's trace how the shape of the data changes as \n",
    "\n",
    "Recall that the input data has dimensions of `torch.Size([64, 1, 28, 28])`. That is, a batch of 64 images with one (greyscale) layer 28x28 pixels each.\n",
    "\n",
    "After the first convolution `Conv2d(1, 10, kernel_size=5)` out comes the tensor of `torch.Size([64, 10, 24, 24])`: stride of 1, padding of 0, and 5x5 kernel produce the result of size `28 - 5 + 1 = 24` in height and width.\n",
    "\n",
    "2d max pooling with 2x2 kernel shrinks the data to half its size in both width and height to `torch.Size([64, 10, 12, 12])`, as `int(dim / kernel_size) = 24/2 = 12`.\n",
    "\n",
    "At the next layer, we apply 5x5 convolution with 20 output layers, `Conv2d(10, 20, kernel_size=5)`. It yields `torch.Size([64, 20, 8, 8])`, sinc `12 - 5 + 1 = 8`.\n",
    "\n",
    "Finally, 2d max pooling with 2x2 kernel halves the the height and width of the input tensor again, making it `torch.Size([64, 20, 4, 4])`. That's a batch of 64 tensors, each having size of 4x4 and with 20 layers, or 320 elements in total (`4 * 4 * 20 = 320`).\n",
    "\n",
    "Now `x.view(-1, 320)` just flattens this tensor into a vector, preserving the first dimension (i.e. the batch size), and we end up with `torch.Size([64, 320])` that can be fed into our first linear layer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.6.2 Linear unit\n",
    "\n",
    "`Linear` is the most fundamental building block of every neural net. Let's create a linear transformation, e.g., from 5 to 2 features and see how it works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Linear(in_features=5, out_features=2, bias=True)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lin = nn.Linear(5, 2)\n",
    "lin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([1., 1., 1., 1., 1.])"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.ones(5)\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.1461,  0.1324], grad_fn=<ThAddBackward>)"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lin(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The results are random because the internal parameters of the `Linear` object were initialized at random:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "(params, bias) = list(lin.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[ 0.0395, -0.0573,  0.2008, -0.0771, -0.2772],\n",
       "        [ 0.1340, -0.2388, -0.2594,  0.2695, -0.0709]], requires_grad=True)"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([0.0252, 0.2980], requires_grad=True)"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bias"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " On a forward pass, `Linear` is just a linear transformation of the input signal (plus the bias):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.1461,  0.1324], grad_fn=<ThAddBackward>)"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "params.matmul(x) + bias"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A linear unit is usually followed by an activation function. In our MNIST example it is ReLU, again.\n",
    "\n",
    "Having two linear layers and a dropout in between allows our model to mix and match the patterns learned in the convolutional layers, and map them to one of 10 digits in the output (provided as labels for training)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.7 Softmax\n",
    "\n",
    "The final activation function of our MNIST model is [softmax](https://en.wikipedia.org/wiki/Softmax_function):\n",
    "\n",
    "```python\n",
    "F.log_softmax(x, dim=1)\n",
    "```\n",
    "\n",
    "It computes a softmax function over the given dimension `dim` of the tensor, followed by a natural logarithm. We have to specify `dim` to compute softmax for each element of the input batch individually.\n",
    "\n",
    "In the MNIST example, ten elements of the softmax output correspond to ten digits, `0` to `9`. *i*-th element of the softmax vector thus represents a probability that the input image corresponds to the digit *i*.\n",
    "\n",
    "A typical softmax output for digit `8` can look like, e.g."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sum = tensor(1.)\n"
     ]
    }
   ],
   "source": [
    "x = torch.tensor([0.0078, 0.0207, 0.0087, 0.0031, 0.0059, 0.0209, 0.0149, 0.0141, 0.8926, 0.0113])\n",
    "print(\"sum =\", x.sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or, in its log form,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-4.8536, -3.8776, -4.7444, -5.7764, -5.1328, -3.8680, -4.2064, -4.2616,\n",
       "        -0.1136, -4.4830])"
      ]
     },
     "execution_count": 88,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.log(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can recover the digit by simply finding the max element:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(8)"
      ]
     },
     "execution_count": 98,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(_, index) = x.max(dim=0)\n",
    "index"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.8 Conclusion\n",
    "\n",
    "In this notebook we've followed the data as it flows through our MNIST model, and learned how the input 28x28 greyscale image gets transformed into one of 10 digits of the target label. In [Part 3](3_mnist_backprop.ipynb) we'll follow the same way *backwards*, propagating the gradient and updating the model parameters."
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}