Worksheet 5, Pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "this worksheets is part of the [mlvu machine learning course](https://mlvu.github.io)<br>\n",
    "setting up your environment: https://bit.ly/3bzpn5C\n",
    "\n",
    "For this worksheet, we'll look into the details of how a deep learning framework operates under the hood. **Tensorflow** and **Pytorch** are both good candidates, but things are a little easier to explain in Pytorch, so we'll use that. Step one, installation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install torch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the following cell executes without error, the installation was succesful."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Worksheet 5: Deep Learning with Pytorch\n",
    "\n",
    "This worksheet assumes that you've watched the _Deep Learning 1_ and _Deep Learning 2_ lectures. If you don't know what backpropation means or what a computation graph is, you should probably give the lectures [a watch](https://mlvu.github.io/) first. If you want an even deeper understanding, try the lectures of the [DLVU course](https://dlvu.github.io), especially the second lecture.\n",
    "\n",
    "The basic datastructure around which all of Pytorch is built is the _Tensor_. It works pretty much exactly like the tensors we've already seen in numpy.\n",
    "\n",
    "To make a matrix filled with given values:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.tensor([[1.0, 2.0], [3.0, 4.0]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To make a matrix filled with (normally distributed) random numbers:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.randn(3, 4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For higher-dimensional matrices (aka tensors), just add more dimensions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.randn(2, 3, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can request the size of a tensor with the ```size()``` function. For the size of a particular dimension, just add the index of that dimension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.randn(3, 4, 2)\n",
    "\n",
    "print(x.size())\n",
    "print(x.size(0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Just like in numpy, we can sum and multiply, and apply basic functions element-wise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a, b = torch.randn(3, 4), torch.rand(3, 4)\n",
    "c = torch.randn(4, 3)\n",
    "\n",
    "# addition, multiplication, etc are all element-wise\n",
    "summed = a + b\n",
    "mult   = a * b\n",
    "power  = a ** 2\n",
    "sine   = torch.sin(a)\n",
    "\n",
    "# _matrix_ multiplication is done through torch.mm\n",
    "mmult = torch.mm(a, c)\n",
    "\n",
    "# Note that the following lines would both give an error. Why?\n",
    "# mult = a * c\n",
    "# torch.mm(a, b)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What are the shapes of the tensors above (summed, mult, etc)? Print them to see if your intuition is correct."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Indexing and slicing works as it does in numpy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(a)\n",
    "\n",
    "print()\n",
    "print('element at the first row and third column of a:', a[0, 2]) # note: zero-indexing\n",
    "print('   first row of a:', a[0, :])\n",
    "print('first column of a:', a[:, 0])\n",
    "\n",
    "# You can also use basic slicing syntax; i:j refers to the range from i to j\n",
    "# (more precisely, i is the first element included, j is the first element \n",
    "#  excluded)\n",
    "print()\n",
    "print('middle two elements of each row of a:\\n', a[:, 1:3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## tensor.view()\n",
    "\n",
    "\n",
    "<div style=\"float:right; width: 250px\">\n",
    "<img src=\"https://upload.wikimedia.org/wikipedia/commons/4/4d/Row_and_column_major_order.svg\" width=\"200px\" title=\"row-major and column-major ordering\" />\n",
    "<small>By Cmglee - Own work, CC BY-SA 4.0, https://commons.wikimedia.org/w/index.php?curid=65107030</small>\n",
    "</div>\n",
    "\n",
    "One of the most important skills in programming in Pytorch is reshaping a tensor. It sounds simple, but there are some important subtleties to pay attention to. If you're not interested in the finer details of Pytorch, you can safely skip to the **backpropagation** section and save this part for later.\n",
    "\n",
    "Computer memory is one big row of bits. However your numbers are arranged in your tensor, in memory, they have to be a  single sequence of of numbers. That means that if you have a matrix like\n",
    "$$\n",
    "\\begin{pmatrix}\n",
    "3 & 5 & 2\\\\\n",
    "1 & 3 & 4\\\\\n",
    "\\end{pmatrix}\n",
    "$$\n",
    "it will be stored in memory as\n",
    "$$\n",
    "\\begin{pmatrix}\n",
    "3 & 5 & 2 & 1 & 3 & 4 \\\\\n",
    "\\end{pmatrix}\\text{.}\n",
    "$$\n",
    "This is called **[row-major ordering](https://en.wikipedia.org/wiki/Row-_and_column-major_order)**: to get the memory layout of the matrix we scan first along the rows, and then along the columns.\n",
    "\n",
    "Pytorch also stores the matrix dimensions, so it can compute that to get element $(2,1)$ in the matrix, it needs to access element 3 in the list.\n",
    "\n",
    "For higher order tensors the principle is the same: the memory layout scans first over the rightmost dimension. So, if we have a tensor $A$ with size ```(2, 2, 3)```, the elements are stored in the order:\n",
    "$$\n",
    "A_{111}, A_{112}, A_{113}, A_{121}, A_{122}, A_{123}, A_{211}, A_{212}, \\ldots\n",
    "$$\n",
    "Note that at every step the rightmost index increments first. When it gets to its maximum, the one to the left of it increments, and so on.\n",
    "\n",
    "### reshaping\n",
    "\n",
    "We can take the data of one matrix in memory, and create a second matrix from it with another shape. This is done by the ```view()``` function, it takes a tensor and gives you a new _view_ on the same data, assuming a different shape:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "matrix = torch.tensor([[3,5,2,6],[1, 3, 4,0],[-1,-3, -4,-0]])\n",
    "\n",
    "print(matrix) \n",
    "print(matrix.view(4, 3))\n",
    "print(matrix.view(2, 6))\n",
    "print(matrix.view(size=(12,)))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can use ```-1``` for one of the arguments. Pytorch will work out what the size of that dimension is from the rest of the values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(matrix.view(-1, 6))\n",
    "print(matrix.view(-1,))\n",
    "# print(matrix.view(-1, 3, -1)) # this doesn't work\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the difference between matrix reshaping and the matrix transpose (done with the ```.t()``` function):\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "matrix = torch.tensor([[3,5,2],[1, 3, 4]])\n",
    "\n",
    "print(matrix.view(3, 2))\n",
    "print(matrix.t())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Transposing can _also_ be done cheaply. Pytorch creates a wrapper around the old matrix that remaps the dimensions. Requesting element $(i, j)$ in the new matrix is remapped to a request for $(j, i)$ in the old matrix. However, in this case, the new matrix is not **contiguous** anymore: the sequence in memory is no longer in row-major order for the new shape. Pytorch can do some things with non-contiguous matrics, but not all. Calling some functions on a non-contiguous matrix will cause an error. \n",
    "\n",
    "For instance, if you try to reshape a non-contiguous matrix with ```view()```, pytorch will complain. The solution is to copy and rearrange the original memory to a new sequence in memory that is contiguous for this matrix shape.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# matrix.t().view(2, 3)              # this fails\n",
    "matrix.t().contiguous().view(2, 3)   # this works, but copies the matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When you use ```reshape()``` instead of view, the matrix is made contiguous if that is necessary, but otherwise the same data is used. This is nice if you don't care much about memory use. If you want to be sure that you're not accidentally copying a large matrix, you should use ```view()```, and check for errors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "matrix.t().reshape(2, 3) # this works, but copies the matrix without warning you"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example: reshaping back and forth\n",
    "\n",
    "Why is this important? Let's look at an example. Let's say you have a large dataset of N images, each with 3 color channels (RGB values) and WxH pixels. You can represent this in a big 4 tensor:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N, C, H, W = 100, 3, 32, 32\n",
    "images = torch.rand(N, H, W, C) # just random values for the example\n",
    "print(images.size())\n",
    "\n",
    "# -- The normal layout for images in pytorch if channels-first, but \n",
    "#    we'll put the channels at the endhere to simplify the example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's say that you want to normalize the data _by color_. That is, we want to rescale the data so that the maximum value in each color channel is 1. We can find the maximum value in a particular dimension, with pytorch (using ```tensor.max(dim=x)```), but to do this over three dimensions requires taking the max three times. \n",
    "\n",
    "An alternative approach is to _reshape_ our data into a matrix, then normalize, and reshape back to the original data.\n",
    "\n",
    "The main thing to remember in this sort of scenario is that we can safely reshape dimensions **that are next to each other**. Here, we want to combine the N, H and W dimensions, and keep the C dimension separate. That means that if we call"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images = images.view(N*H*W, C)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```images``` is now a _matrix_, with the N, H and W dimensions flattened (in row-major order) along the  vertical dimension. We can now easily normalize this matrix vertically:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images = images / images.max(dim=0)[0]\n",
    "\n",
    "# -- Note that even though the numerator has size (102400, 3) and \n",
    "#    the denominator has size(3), we can still do element-wise division,\n",
    "#    because pytorch supports broadcasting, the same as numpy. It's \n",
    "#    worth knowing _exactly_ how broadcasting works, or the results\n",
    "#    can be surprising. Check the numpy worksheet for more information."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that all the colors are normalized, we can simply reshape our data back into images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images = images.view(N, H, W, C)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This sort of thing can be very useful, but you have to be careful to keep in mind what the layout of the data is in memory. Pytorch happily let you do"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images = images.view(N*H, W*C)\n",
    "images = images.view(H, N, C, W)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "but the result is is basically a completely scrambled dataset. \n",
    "\n",
    "In short, make sure you understand the difference between _transposing_ (swapping dimensions) and reshaping (changing the dimensions, but keeping the data the same)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Backpropagation\n",
    "\n",
    "Let's look at how Pytorch implements _backpropagation_. All we need to do is create some tensors, tell Pytorch that we want to compute gradients for them, and then do some operations on them that result in a single scalar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a, b = torch.randn(3, 4), torch.rand(3, 4)   # create some tensors ...\n",
    "a.requires_grad = True   # ... tell pytorch that we want gradients on a ...\n",
    "\n",
    "out = ((a + b) ** 2).sum()   # ... and perform a computation resulting in a scalar value.\n",
    "print(out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Pytorch has not just computed the result (```out```), it's also included a pointer to a ```SumBackward``` object representing the computation (summing) that created ```out```. This object links back to other objects, all the way down to the start of the computation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(out.grad_fn) # sum\n",
    "print(out.grad_fn.next_functions[0][0]) # raise to a power\n",
    "print(out.grad_fn.next_functions[0][0].next_functions[0][0]) # addition\n",
    "\n",
    "# Note: these are not attributes you would normally use. We just call them here to show that \n",
    "# Pytorch is remembering everything we do.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We've asked Pytorch to ensure we can compute a gradient on ```a``` and done some basic computation. The computation has resulted in a single number (```out```), so we can now compute the gradient of ```a``` over that output. Remember, backpropagation only works efficiently if the output of the computation graph is a single scalar, usually your loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(a.grad)     # this is the gradient on a. Note it's currently empty\n",
    "\n",
    "out.backward()    # ask Pytorch to perform backprop\n",
    "\n",
    "print(a.grad)     # now a has a gradient"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the gradient of ```a``` has the same shape as ```a```."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning\n",
    "\n",
    "Pytorch has many utilities to help you quickly build elaborate networks, but it's instructive to first see how you would use just these tools to build a simple model. As an example, we will build a simple linear regression model.\n",
    "\n",
    "First, let's generate some simple random data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.randn(1000, 32)                      # 1000 instances, with 32 features\n",
    "wt, bt = torch.randn(32, 1), torch.randn(1)    # function to compute the true labels\n",
    "t = torch.mm(x, wt) + bt                       # the true target labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next up, we define the parameters of our model (we'll initialize them randomly)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "w = torch.randn(32, 1, requires_grad=True) \n",
    "b = torch.randn(1, requires_grad=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that any method that creates tensors (like ```torch.randn()```) can be told that it should make them require a gradient).\n",
    "\n",
    "Here's what one computation of the model output over the whole data looks like. We'll print the shapes of the tensors to see what's going on. **Before you run this cell, see if you can work out what the sizes will be.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print('          data size:', x.size())\n",
    "\n",
    "# model output\n",
    "y = torch.mm(x, w) + b\n",
    "\n",
    "print('        output size:', y.size())\n",
    "\n",
    "print()\n",
    "print('first 3 predictions:', y[:3, 0])\n",
    "print('       ground truth:', t[:3, 0]) # note that these will be completely different, because\n",
    "                                        # we haven't started training yet\n",
    "\n",
    "# residuals\n",
    "r = t - y\n",
    "print()\n",
    "print('     residuals size:', r.size())\n",
    "    \n",
    "# mean-squared-error loss \n",
    "loss = (r ** 2).mean()\n",
    "print()\n",
    "print('               loss:', loss.item())\n",
    "# -- if you have a tensor with a single number, .item() will turn it into a normal float for you.\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now apply backpropagation, and see that we get a gradient over our two parameters ```w``` and ```b```. Before you run the cell, what will the sizes of the gradient tensors be?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss.backward()\n",
    "\n",
    "print('gradient on w:', w.grad)\n",
    "print('gradient on b:', b.grad)\n",
    "\n",
    "# NB: if you run the cell twice, pytorch will complain. After each backward, pytorch expects a new forward."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now ready to build a training loop. We'll use basic gradient descent without minibatches, computing the loss over the whole data every iteration. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hyperparameters\n",
    "iterations = 21\n",
    "learning_rate= 0.5\n",
    "\n",
    "# regenerate the data and model\n",
    "x = torch.randn(1000, 32)                      # 1000 instances, with 32 features\n",
    "wt, bt = torch.randn(32, 1), torch.randn(1)    # parameters of the true model\n",
    "t = torch.mm(x, wt) + bt \n",
    "\n",
    "w = torch.randn(32, 1, requires_grad=True) \n",
    "b = torch.randn(1, requires_grad=True)\n",
    "\n",
    "for i in range(iterations):\n",
    "\n",
    "    # forward pass\n",
    "    y = torch.mm(x, w) + b\n",
    "\n",
    "    # mean-squared-error loss \n",
    "    r = t - y\n",
    "    loss = (r ** 2).mean()\n",
    "\n",
    "    # backpropagation\n",
    "    loss.backward()\n",
    "        \n",
    "    # print the loss\n",
    "    print(f'iteration {i: 4}: loss {loss:.4}')\n",
    "    \n",
    "    # gradient descent\n",
    "    w.data = w.data - learning_rate * w.grad.data\n",
    "    b.data = b.data - learning_rate * b.grad.data\n",
    "    # -- Note that we don't want the autodiff engine to compute gradients over this part.\n",
    "    #   by operrating on w.data, we are only changing the values of the tensor not \n",
    "    #   remembering a computation graph.\n",
    "\n",
    "    # delete the gradients\n",
    "    w.grad.data.zero_()\n",
    "    b.grad.data.zero_()\n",
    "    # -- if we don't do this, the gradients are remembered, and any new gradients are added\n",
    "    #    to the old.   \n",
    "\n",
    "# show the true model, and the learned model\n",
    "print()\n",
    "print('true model: ', wt.data[:4].t(), bt.data)\n",
    "print('learned model:', w.data[:4].t(), b.data)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## torch.optim\n",
    "\n",
    "While you can build everything yourself using these low-level objects, many people have done so before and packaged the results in reusable libraries. \n",
    "\n",
    "We'll first look at ```torch.optim```, which contains a number of _optimizers_. Using these, we don't have to implement the gradient descent step ourselves. This may not seem like a big part of our code, but it can get more complicated when we want to try variations on gradient descent like Adam. To illustrate, let's use the Adam optimizer in our linear regression example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.optim import Adam\n",
    "\n",
    "# hyperparameters\n",
    "iterations = 101\n",
    "learning_rate= 2.0\n",
    "\n",
    "# regenerate the data and model\n",
    "x = torch.randn(1000, 32)                      # 1000 instances, with 32 features\n",
    "wt, bt = torch.randn(32, 1), torch.randn(1)    # parameters of the true model\n",
    "t = torch.mm(x, wt) + bt \n",
    "\n",
    "w = torch.randn(32, 1, requires_grad=True) \n",
    "b = torch.randn(1, requires_grad=True)\n",
    "\n",
    "# Create the optimizer. It needs to know two things:\n",
    "# - the learning rate\n",
    "# - which parameters its responsible for\n",
    "opt = Adam(lr=learning_rate, params=[w, b])\n",
    "\n",
    "for i in range(iterations):\n",
    "\n",
    "    # forward/backward, same as before\n",
    "    y = torch.mm(x, w) + b\n",
    "    r = t - y\n",
    "    loss = (r ** 2).mean()\n",
    "    \n",
    "    loss.backward() \n",
    "    # -- Note that the optimizer _doesn't_ compute the gradients. We still\n",
    "    #    do that ourselves. The optimizer takes the gradients, and uses them\n",
    "    #    to adapt the parameters. \n",
    "        \n",
    "    # print the loss\n",
    "    if i % 20 == 0:\n",
    "        print(f'iteration {i: 4}: loss {loss:.4}')\n",
    "    \n",
    "    # Perform the gradient descent step\n",
    "    opt.step() \n",
    "    \n",
    "    # The optimizer can zero the gradients for us \n",
    "    # (but we still have to tell it to do so)\n",
    "    opt.zero_grad()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You may have noticed that Adam (which is supposed to be better than gradient descent in many ways), actually takes longer to converge. That's because this is a very simple problem. Adam's strength shows when you train very large networks with many weights all doing different things."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## torch.nn\n",
    "\n",
    "The package ```torch.nn``` contains utilities for building the kind of large, complex neural networks we've seen in the lectures. It is built around <em>modules</em>: modules are classes that combine a set of weights with a ```forward()``` function that uses these weights to compute a forward pass.\n",
    "\n",
    "The simplest module is probably ```torch.nn.Linear```. It implements a simple linear operation with a weight matrix and a bias vector (this is like the ```Dense``` layer in Keras)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "\n",
    "lin = nn.Linear(2, 2) # a linear function from a 2-vector to a 2-vector\n",
    "\n",
    "print('weight matrix:', lin.weight)\n",
    "print()\n",
    "print('bias vector:', lin.bias)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see that Pytorch knows that ```lin.weight``` and ```lin.bias``` are _parameters_. This is because they are ```nn.Parameter``` objects, a lightweight wrapper around the torch tensor, which signals that this tensor is meant to be treated like a model parameter. It has ```requires_grad=True``` by default, and there are helper functions to collect all parameters of a complex model.\n",
    "\n",
    "We can apply the linear transformation by calling ```lin``` just like a function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.tensor([1.0, 2.0])\n",
    "lin(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the resulting tensor has a ```grad_fn``` attribute, so we can tell that the computation graph is being remembered.\n",
    "\n",
    "To implement a module of our own, we create a subclass of the ```nn.Module``` class. All we need to implement is the constructor and the ```forward``` function. Here is a module for a simple two-layer MLP with a ReLU activation on the hidden layer.\n",
    "\n",
    "To illustrate how to define and apply parameters, we will also add a multiplier to the output (a single learnable value). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F # some helpful utility functions\n",
    "\n",
    "class MLP(nn.Module):\n",
    "    \n",
    "    def __init__(self, in_size = 16, hidden_size=32, out_size=1):\n",
    "        \"\"\"\n",
    "        This is the _constructor_ the function that creates an instance of the MLP class.\n",
    "        \n",
    "        The argument 'self' is a standard argument in python object-oriented programming. It\n",
    "        refers to the current instance that we're creating. The other arguments are parameters\n",
    "        of the MLP.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        \n",
    "        # everything that has parameters should be created in the contructor\n",
    "        self.layer1 = nn.Linear(in_size, hidden_size)\n",
    "        self.layer2 = nn.Linear(hidden_size, out_size)\n",
    "        \n",
    "        # the layers have most of the parameters, but we will also add one of our own\n",
    "        self.mult = nn.Parameter(torch.tensor([1.0]))\n",
    "        # -- we create a tensor with the initial value, and wrap it in an nn.Parameter\n",
    "        #    objects. Because it's an nn.Parameter, pytorch will take care of the rest.\n",
    "        \n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        This is the function that gets executed when we call the module like a  function.\n",
    "\n",
    "        - The argument 'self' again refers to the current object.\n",
    "        - The argument 'x' is the input to the function (multiple arguments, named aguments \n",
    "          and even no arguments are possible)\n",
    "        \"\"\"        \n",
    "\n",
    "        h = self.layer1(x)  # apply the first layer\n",
    "        h = F.relu(h)       # apply a ReLU nonlinearity\n",
    "        o = self.layer2(h)  # apply the second layer    \n",
    "        \n",
    "        o = o * self.mult        # apply the multiplier\n",
    "\n",
    "        return o"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now create an MLP instance, and feed it some data. Before you run the cell, can you predict the size of the output?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlp = MLP()            # create an MLP with the standard dimensions\n",
    "\n",
    "x = torch.randn(3, 16) # three instances, with 16 features\n",
    "\n",
    "mlp(x)                 # pass the data through the MLP."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because we've subclassed ```nn.Module```, we get a lot of functionality for free. For instance, ```mlp``` has a function that lets us loop over all its parameters and the parameters of its modules (and the parameters of their modules and so on)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for param in mlp.parameters():\n",
    "    print(param.size())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order, these are the multiplier, the weights matrix of the first layer, the bias of the first layer, the weight matrix of the second layer and the bias of the second layer.\n",
    "\n",
    "This is helpful when we need to let the optimizer know what the parameters of our model are. \n",
    "\n",
    "Here's an example of how to put everything together and train our MLP on some generated data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hyperparameters\n",
    "iterations = 1000\n",
    "learning_rate= 0.01\n",
    "\n",
    "# regenerate the data and model\n",
    "x = torch.randn(1000, 32)                          # 1000 instances, with 2 features\n",
    "t = torch.sqrt(x.pow(2).sum(dim=1, keepdim=True))  # we'll use the vector norm as a target function\n",
    "\n",
    "model = MLP(32, 64, 1)\n",
    "\n",
    "opt = Adam(lr=learning_rate, params=model.parameters())\n",
    "# -- Note that we just point the optimizer to the parameters generator\n",
    "\n",
    "for i in range(iterations):\n",
    "\n",
    "    y = model(x)\n",
    "    loss = F.mse_loss(y, t) \n",
    "    # -- We'll switch to the pytorch implementation of the MSE. \n",
    "    \n",
    "    if i % 50 == 0:\n",
    "        print(f'iteration {i: 4}: loss {loss:.4}')\n",
    "    \n",
    "    loss.backward()\n",
    "    \n",
    "    opt.step()\n",
    "    opt.zero_grad()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Piece de resistance: the variational autoencoder\n",
    "\n",
    "In the Keras worksheet, we built an autoencoder. Now, let's build a _variational_ autoencoder (VAE). This possible in Keras as well, but several aspects of the VAE make it a bit awkward: we have one loss at the end and one loss halfway down the network. Also we have to implement a sampling step in the middle. \n",
    "\n",
    "In pytorch all this becomes a bit easier, because the forward pass of our model looks so much more like regular code.\n",
    "\n",
    "**Note that this is only a tutorial on how to _build_ a VAE. We'll expect you to know how a VAE works already. If you don't, please review the second Deep Learning lecture.**\n",
    "\n",
    "We'll start by importing the MNIST data. First we have to install the ```torchvision``` package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install torchvision"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we load the data. In pytorch, data is usually packaged in a dataloader, which can efificiently serve you batches of data. To simplify things, we'll only load the training data. We will ask for batches of 32 images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision\n",
    "from torchvision import transforms\n",
    "\n",
    "train = torchvision.datasets.MNIST(root='./mnist', train=True, download=True, transform=transforms.ToTensor())\n",
    "trainloader = torch.utils.data.DataLoader(train, batch_size=32, shuffle=True, num_workers=2)\n",
    "\n",
    "# We use the trainloader by looping over it. In this case, each batch is a pair of an image tensor, and a \n",
    "# label tensor. We'll ignore the labels.\n",
    "for images, labels in trainloader:\n",
    "    print(images.size(), labels.size())\n",
    "    break\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next up, we'll define the model. As we did in the Keras worksheet, we'll define the encoder and decoder separately"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pytorch doesn't give us a Reshape module. We'll add that ourselves so we \n",
    "# can define the encoder and decoder as sequences of operations\n",
    "class Reshape(nn.Module):\n",
    "    def __init__(self, shape):\n",
    "        super().__init__()\n",
    "        self.shape = shape\n",
    "\n",
    "    def forward(self, input):\n",
    "        return input.view( (input.size(0),) + self.shape) # keep the batch dimensions, reshape the rest\n",
    "\n",
    "# - channel sizes\n",
    "a, b, c = 16, 32, 128\n",
    "latent_size = 2\n",
    "\n",
    "encoder = nn.Sequential(\n",
    "    nn.Conv2d(1, a, (3, 3), padding=1), nn.ReLU(),\n",
    "    nn.MaxPool2d((2, 2)),\n",
    "    nn.Conv2d(a, b, (3, 3), padding=1), nn.ReLU(),\n",
    "    nn.Conv2d(b, b, (3, 3), padding=1), nn.ReLU(),    \n",
    "    nn.MaxPool2d((2, 2)),\n",
    "    nn.Conv2d(b, c, (3, 3), padding=1), nn.ReLU(),\n",
    "    nn.Conv2d(c, c, (3, 3), padding=1), nn.ReLU(),\n",
    "    nn.MaxPool2d((2, 2)),\n",
    "    nn.Flatten(),\n",
    "    nn.Linear(3 * 3 * c, 2 * latent_size)\n",
    ")\n",
    "\n",
    "decoder = nn.Sequential(\n",
    "    nn.Linear(latent_size, c * 3 * 3), nn.ReLU(),\n",
    "    Reshape((c, 3, 3)),\n",
    "    nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),\n",
    "    nn.ConvTranspose2d(c, b, (3, 3), padding=1), nn.ReLU(),\n",
    "    nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),\n",
    "    nn.ConvTranspose2d(b, a, (3, 3), padding=0), nn.ReLU(), # note the padding\n",
    "    nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),\n",
    "    nn.ConvTranspose2d(a, 1, (3, 3), padding=1), nn.Sigmoid()\n",
    ")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are cheating a little with the output activation. To truly follow the derivation of the VAE, this should define a distribution in the data space (which are continuous numbers _between_ 0 and 1). In this case, we use a sigmoid activation with a binary cross-entropy, which would be a distribution for binary data (either 0 or 1). \n",
    "\n",
    "A more correct VAE would use, for instance, a Gaussian output, which boils down to an MSE loss (but doesn't work well for this task), or add some terms to make the BCE loss work theoretically. For now we'll ignore this and stick with a plain BCE loss.\n",
    "\n",
    "More importantly, note that the output of the encoder is _twice_ the size of the latent space, while the input to the decoder is the size of the latent space. This is because the decoder gives us a mean _and a variance_ on the latent space, from which we'll _sample_ the input to the decoder.\n",
    "\n",
    "We'll also need to compute the KL loss term from this mean and variance. We'll introduce some utility functions for both operations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def kl_loss(zmean, zsig):\n",
    "    b, l = zmean.size()\n",
    "\n",
    "    kl = 0.5 * torch.sum(zsig.exp() - zsig + zmean.pow(2) - 1, dim=1)\n",
    "    # -- The KL divergence between a given normal distribution and a standard normal distribution\n",
    "    #    can be rewritten this way. It's a good exercise to work this out.\n",
    "\n",
    "    assert kl.size() == (b,)\n",
    "    # -- At this point we want the loss to be a single value of each instance in the batch.\n",
    "    #    Asserts like this are a good way to document what you know about the shape of the \n",
    "    #    tensors you deal with.\n",
    "\n",
    "    return kl\n",
    "\n",
    "def sample(zmean, zsig):\n",
    "    b, l = zmean.size()\n",
    "\n",
    "    # sample epsilon from a standard normal distribution\n",
    "    eps = torch.randn(b, l)\n",
    "\n",
    "    # transform eps to a sample from the given distribution\n",
    "    return zmean + eps * (zsig * 0.5).exp()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want to try to work out why the KL term looks like that, you should [start here](https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence#Multivariate_normal_distributions) and try to rewrite to pytorch code.\n",
    "\n",
    "Now that we have all the steps in our model, we can build the training loop.\n",
    "\n",
    "Since this might take a while, we install and import [```tqdm```](https://tqdm.github.io/) to give us some nice progress bars."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "parameters = list(encoder.parameters()) + list(decoder.parameters()) # -- retrieve all paremeters of both models\n",
    "opt = Adam(lr=0.003, params=parameters)\n",
    "\n",
    "for epoch in range(5):\n",
    "    for images, _ in tqdm(trainloader): # if tqdm gives you trouble just remove it\n",
    "        b, c, h, w = images.size()\n",
    "                \n",
    "        # forward pass\n",
    "        z = encoder(images) \n",
    "        \n",
    "        # - split z into mean and sigma\n",
    "        zmean, zsig = z[:, :latent_size], z[:, latent_size:]\n",
    "        kl = kl_loss(zmean, zsig)\n",
    "        \n",
    "        zsample = sample(zmean, zsig)\n",
    "        \n",
    "        o = decoder(zsample)\n",
    "        rec = F.binary_cross_entropy(o, images, reduction='none')\n",
    "        rec = rec.view(b, c*h*w).sum(dim=1)\n",
    "        # -- Reconstruction loss. We ask pytorch not to sum the loss, and sum over the \n",
    "        #    channels and pixels ourselves. This gives us a loss per instance that we \n",
    "        #    can add to the kl loss\n",
    "        \n",
    "        loss = (rec + kl).mean() # sum the losses and take the mean\n",
    "        loss.backward()\n",
    "        \n",
    "        opt.step()\n",
    "        opt.zero_grad()\n",
    "    \n",
    "    print(f'epoch {epoch}: {loss.item()}')\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Training takes about 2 minutes per epoch on my laptop (note that the autoencoder is a bit bigger than the one in the Keras worksheet). I recommend letting it run for at least 3 epochs.\n",
    "\n",
    "To visualize what our autoencoder did, we'll plot some of the data by their latent space coordinates. That is, for each image, we compute the latent space coordinates, and plot the original image at that point (this takes a while):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import matplotlib as mpl\n",
    "import math\n",
    "import numpy as np\n",
    "\n",
    "# gather up first 200 batches into one big tensor\n",
    "numbatches = 200 # -- change 200 to a lower number to speed things up\n",
    "images, labels = [], []\n",
    "for i, (ims, lbs) in enumerate(trainloader):\n",
    "    images.append(ims)\n",
    "    labels.append(lbs)\n",
    "    \n",
    "    if i > numbatches:\n",
    "        break\n",
    "    \n",
    "images, labels = torch.cat(images, dim=0), torch.cat(labels, dim=0)\n",
    "\n",
    "n, c, h, w = images.size()\n",
    "\n",
    "z = encoder(images)\n",
    "latents = z[:, :2].data\n",
    "\n",
    "mn, mx = latents.min(), latents.max()\n",
    "size = 1.0 * (mx-mn)/math.sqrt(n)\n",
    "# Change 0.75 to any value between ~ 0.5 and 1.5 to make the digits smaller or bigger\n",
    "\n",
    "fig = plt.figure(figsize=(16,16))\n",
    "\n",
    "# colormap for the images\n",
    "norm = mpl.colors.Normalize(vmin=0,vmax=9)\n",
    "cmap = mpl.cm.get_cmap('tab10')\n",
    "\n",
    "for i in range(n):\n",
    "    \n",
    "    x, y = latents[i, 0:2]\n",
    "    l = labels[i]\n",
    "    \n",
    "    im = images[i, :]\n",
    "    alpha_im = im.permute(1, 2, 0).numpy()\n",
    "    color = cmap(norm(l))\n",
    "    color_im = np.asarray(color)[None, None, :3]\n",
    "    color_im = np.broadcast_to(color_im, (h, w, 3))\n",
    "    # -- To make the digits transparent we make them solid color images and use the \n",
    "    #    actual data as an alpha channel.\n",
    "    #    color_im: 3-channel color image, with solid color corresponding to class\n",
    "    #    alpha_im: 1-channel grayscale image corresponding to input data\n",
    "\n",
    "    im = np.concatenate([color_im, alpha_im], axis=2)\n",
    "    plt.imshow(im, extent=(x, x + size, y, y + size))\n",
    "\n",
    "    plt.xlim(mn, mx)\n",
    "    plt.ylim(mn, mx)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that even with just 2 latent dimensions, the VAE manages to separate out the digits quite well (remember, the model doesn't see the labels, they are only used to color the plot). Even within a digit's, the variations (like angle, and stroke thickness) are being grouped together.\n",
    "\n",
    "You can try to make the separation clearer by training for more epochs, by adding more channels or layers to the encoder or decoder, by tuning the learning rate more precisely.\n",
    "\n",
    "## Conclusion\n",
    "\n",
    "That's it, you've finished the last worksheet. Pytorch is a complicated system, so don't feel too bad if you don't understand all the details. Play around with the code, and follow some of the links below, and you'll soon deepen your understanding.\n",
    "\n",
    "### Further reading\n",
    "\n",
    "* Pytorch 60 minute blitz: https://pytorch.org/tutorials/beginner/deep_learning_60min_blitz.html\n",
    "* Learning Pytorch with examples: https://pytorch.org/tutorials/beginner/pytorch_with_examples.html\n",
    "* Visualizing Models, Data, and Training with TensorBoard: https://pytorch.org/tutorials/intermediate/tensorboard_tutorial.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}