Lab1_GeometricDeepLearning_Introduction_to_PyTorch.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Lab1_GeometricDeepLearning_Introduction_to_PyTorch.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "wSXfdKWTmJpv",
        "XafS-KDNmJpy",
        "epWXbjpYmJp0",
        "hJlASM8fmJp3",
        "L1jb4kYzmJp4",
        "eZng3I-1mJp5",
        "xPfGyYUtmJp8",
        "mjS31RRgmJp9",
        "3oWriws6mJp_",
        "Ggwnur38mJqA",
        "sfIXs_b5mJqB",
        "RtixEHqqmJqB",
        "dtqyEhuUmJqC",
        "Sh-F2kS6mJqE",
        "osOyNndimJqF",
        "sTAqvA0QmJqG"
      ]
    },
    "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.7.7"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xvrobO8FmJpd"
      },
      "source": [
        "# ACM Summer School 2021\n",
        "\n",
        "## Lab : Geometric DeepLearning - Introduction to PyTorch\n",
        "## Presenter : Aditya Intwala\n",
        "\n",
        "![Status](https://img.shields.io/static/v1.svg?label=Status&message=Finished&color=green)\n",
        "\n",
        "**Exercise notebook:** \n",
        "[![View on Github](https://img.shields.io/static/v1.svg?logo=github&label=Repo&message=View%20On%20Github&color=lightgrey)](https://github.com/adityaintwala/ACM-GeometricDeepLearning-2021/blob/main/Lab1_GeometricDeepLearning_Introduction_to_PyTorch.ipynb)\n",
        "[![Open In Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adityaintwala/ACM-GeometricDeepLearning-2021/blob/main/Lab1_GeometricDeepLearning_Introduction_to_PyTorch.ipynb)\n",
        "**Solution notebook:** \n",
        "[![View on Github](https://img.shields.io/static/v1.svg?logo=github&label=Repo&message=View%20On%20Github&color=lightgrey)](https://github.com/adityaintwala/ACM-GeometricDeepLearning-2021/blob/main/Solution/Lab1_GeometricDeepLearning_Introduction_to_PyTorch.ipynb)\n",
        "[![Open In Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adityaintwala/ACM-GeometricDeepLearning-2021/blob/main/Solution/Lab1_GeometricDeepLearning_Introduction_to_PyTorch.ipynb)  \n",
        "\n",
        "*This is an empty version of the notebook for the Lab of Geometric Deep Learning session at ACM Summer School 2021.* \n",
        "\n",
        "*Try to fill in the blanks indicated by `## TODO` comments and `...` parts in the code*"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vJZLbUsYmJpg"
      },
      "source": [
        "Welcome to PyTorch tutorial for the Deep Learning 2021 at the ACM Summer School! The following notebook is meant to give a short introduction to PyTorch basics, and get you setup for writing your own neural networks. PyTorch is an open source machine learning framework that allows you to write your own neural networks and optimize them efficiently. However, PyTorch is not the only framework of its kind. Alternatives to PyTorch include [TensorFlow](https://www.tensorflow.org/), [Keras](https://keras.io/api/) and [Caffe](http://caffe.berkeleyvision.org/).\n",
        "\n",
        "There are many great tutorials online, including the [\"60-min blitz\"](https://pytorch.org/tutorials/beginner/deep_learning_60min_blitz.html) on the official [PyTorch website](https://pytorch.org/tutorials/). \n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Luegvlm9mJpg"
      },
      "source": [
        "## Standard libraries\n",
        "import os\n",
        "import math\n",
        "import numpy as np \n",
        "import time\n",
        "\n",
        "## Imports for plotting\n",
        "import matplotlib.pyplot as plt\n",
        "%matplotlib inline \n",
        "from matplotlib.colors import to_rgba\n",
        "import seaborn as sns\n",
        "sns.set()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fAU-gzW8mJph"
      },
      "source": [
        "## The Basics of PyTorch\n",
        "\n",
        "We will start with reviewing the very basic concepts of PyTorch. As a prerequisite, we recommend to be familiar with the `numpy` package as most machine learning frameworks are based on very similar concepts. If you are not familiar with numpy yet, don't worry: here is a [tutorial](https://numpy.org/devdocs/user/quickstart.html) to go through. \n",
        "\n",
        "So, let's start with importing PyTorch. The package is called `torch`, based on its original framework [Torch](http://torch.ch/). As a first step, we can check its version:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cgWOto4fmJph",
        "outputId": "c1275200-d55c-466d-af9e-fdd0728a5375"
      },
      "source": [
        "import torch\n",
        "print(\"Using torch\", torch.__version__)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Using torch 1.6.0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "43jmP3gTmJpi"
      },
      "source": [
        "As in every machine learning framework, PyTorch provides functions that are stochastic like generating random numbers. However, a very good practice is to setup your code to be reproducible with the exact same random numbers. This is why we set a seed below. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "la8Q8m-0mJpi",
        "outputId": "c17150ce-5c66-4b6a-f790-8f3e33d4525b"
      },
      "source": [
        "torch.manual_seed(7) # Setting the seed"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<torch._C.Generator at 0x1d5ea819910>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 4
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fn1vl-kHmJpi"
      },
      "source": [
        "### Tensors\n",
        "\n",
        "Tensors are the PyTorch equivalent to Numpy arrays, with the addition to also have support for GPU acceleration.\n",
        "The name \"tensor\" is a generalization of concepts you already know. For instance, a vector is a 1-D tensor, and a matrix a 2-D tensor. When working with neural networks, we will use tensors of various shapes and number of dimensions.\n",
        "\n",
        "Most common functions you know from numpy can be used on tensors as well. Actually, since numpy arrays are so similar to tensors, we can convert most tensors to numpy arrays (and back) but we don't need it too often.\n",
        "\n",
        "#### Initialization\n",
        "\n",
        "Let's first start by looking at different ways of creating a tensor. There are many possible options, the most simple one is to call `torch.Tensor` passing the desired shape as input argument:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YgxKIi7tmJpj",
        "outputId": "fe8d5754-b49b-4c7d-c2a5-262feeb1f593"
      },
      "source": [
        "x = torch.Tensor(2, 3, 4)\n",
        "print(x)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "tensor([[[9.9184e-39, 1.0561e-38, 5.9694e-39, 8.9082e-39],\n",
            "         [1.0194e-38, 9.1837e-39, 4.6837e-39, 9.2755e-39],\n",
            "         [1.0837e-38, 8.4490e-39, 9.1836e-39, 1.0653e-38]],\n",
            "\n",
            "        [[8.9082e-39, 9.9184e-39, 9.0000e-39, 1.0561e-38],\n",
            "         [1.0653e-38, 4.1327e-39, 8.9082e-39, 9.8265e-39],\n",
            "         [9.4592e-39, 1.0561e-38, 1.0653e-38, 1.0469e-38]]])\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CoOnc8QSmJpm"
      },
      "source": [
        "The function `torch.Tensor` allocates memory for the desired tensor, but reuses any values that have already been in the memory. To directly assign values to the tensor during initialization, there are many alternatives including:\n",
        "\n",
        "* `torch.zeros`: Creates a tensor filled with zeros\n",
        "* `torch.ones`: Creates a tensor filled with ones\n",
        "* `torch.rand`: Creates a tensor with random values uniformly sampled between 0 and 1\n",
        "* `torch.randn`: Creates a tensor with random values sampled from a normal distribution with mean 0 and variance 1\n",
        "* `torch.arange`: Creates a tensor containing the values $N,N+1,N+2,...,M$\n",
        "* `torch.Tensor` (input list): Creates a tensor from the list elements you provide\n",
        "* `torch.from_numpy`: Converts a numpy array into a PyTorch tensor"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YKTsR1vmmJpn"
      },
      "source": [
        "## TODO: Create a tensor from the list [[25, 4], [50, 2]]\n",
        "x = ...\n",
        "print(x)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qfi2fLEwmJpo"
      },
      "source": [
        "## TODO: Create a tensor with random values between 0 and 1 with the shape [4, 3, 2]\n",
        "x = ...\n",
        "print(x)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wHkJyqrRmJpp"
      },
      "source": [
        "You can obtain the shape of a tensor in the same way as in numpy (`x.shape`), or using the `.size` method:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lWZaW5iamJpq"
      },
      "source": [
        "shape = x.shape\n",
        "print(\"Shape:\", x.shape)\n",
        "\n",
        "size = x.size()\n",
        "print(\"Size:\", size)\n",
        "\n",
        "dim1, dim2, dim3 = x.size()\n",
        "print(\"Size:\", dim1, dim2, dim3)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pv36R_PWmJpr"
      },
      "source": [
        "#### Operations\n",
        "\n",
        "Most operations that exist in numpy, also exist in PyTorch. A full list of operations can be found in the [PyTorch documentation](https://pytorch.org/docs/stable/tensors.html#), but we will review the most important ones here.\n",
        "\n",
        "The simplest operation is to add two tensors:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ku2u5KuAmJps"
      },
      "source": [
        "x1 = torch.rand(2, 3)\n",
        "x2 = torch.rand(2, 3)\n",
        "\n",
        "## TODO: Sum up the tensors x1 and x2, and store the result in y\n",
        "y = ...\n",
        "\n",
        "print(\"X1\", x1)\n",
        "print(\"X2\", x2)\n",
        "print(\"Y\", y)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ytoEPGqSmJps"
      },
      "source": [
        "Calling `x1 + x2` creates a new tensor containing the sum of the two inputs. However, we can also use in-place operations that are applied directly on the memory of a tensor. We therefore change the values of `x2` without the chance to re-accessing the values of `x2` before the operation. An example is shown below:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FlCqfsNnmJps"
      },
      "source": [
        "x1 = torch.rand(2, 3)\n",
        "x2 = torch.rand(2, 3)\n",
        "print(\"X1 (before)\", x1)\n",
        "print(\"X2 (before)\", x2)\n",
        "\n",
        "x2.add_(x1)\n",
        "print(\"X1 (after)\", x1)\n",
        "print(\"X2 (after)\", x2)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DwZwz_iLmJpt"
      },
      "source": [
        "In-place operations are usually marked with a underscore postfix (e.g. \"add_\" instead of \"add\").\n",
        "\n",
        "Another common operation aims at changing the shape of a tensor. A tensor of size (2,3) can be re-organized to any other shape with the same number of elements (e.g. a tensor of size (6), or (3,2), ...). In PyTorch, this operation is called `view`:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ajq4CbbMmJpt"
      },
      "source": [
        "## TODO: Create a tensor with the input values [0, .., 5] (i.e. a range to 6)\n",
        "x = ...\n",
        "print(\"X\", x)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ouGmhDDtmJpt"
      },
      "source": [
        "x = x.view(2, 3)\n",
        "print(\"X\", x)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pZllx0XbmJpt"
      },
      "source": [
        "x = x.permute(1, 0) # Swapping dimension 0 and 1\n",
        "print(\"X\", x)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DLSkLe3wmJpu"
      },
      "source": [
        "Other commonly used operations include matrix multiplications, which are essential for neural networks. Quite often, we have an input vector $\\mathbf{x}$, which is transformed using a learned weight matrix $\\mathbf{W}$. There are multiple ways and functions to perform matrix multiplication, some of which we list below:\n",
        "\n",
        "* `torch.matmul`: Performs the matrix product over two tensors, where the specific behavior depends on the dimensions. If both inputs are matrices (2-dimensional tensors), it performs the standard matrix product. For higher dimensional inputs, the function supports broadcasting (for details see the [documentation](https://pytorch.org/docs/stable/generated/torch.matmul.html?highlight=matmul#torch.matmul)). Can also be written as `a @ b`, similar to numpy. \n",
        "* `torch.mm`: Performs the matrix product over two matrices, but doesn't support broadcasting (see [documentation](https://pytorch.org/docs/stable/generated/torch.mm.html?highlight=torch%20mm#torch.mm))\n",
        "* `torch.bmm`: Performs the matrix product with a support batch dimension. If the first tensor $T$ is of shape ($b\\times n\\times m$), and the second tensor $R$ ($b\\times m\\times p$), the output $O$ is of shape ($b\\times n\\times p$), and has been calculated by performing $b$ matrix multiplications of the submatrices of $T$ and $R$: $O_i = T_i @ R_i$\n",
        "* `torch.einsum`: Performs matrix multiplications and more (i.e. sums of products) using the Einstein summation convention. Explanation of the Einstein sum can be found in assignment 1.\n",
        "\n",
        "Usually, we use `torch.matmul` or `torch.bmm`. We can try a matrix multiplication with `torch.matmul` below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "slxiuKT7mJpu"
      },
      "source": [
        "x = torch.arange(6)\n",
        "x = x.view(2, 3)\n",
        "print(\"X\", x)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uJudqXo-mJpv"
      },
      "source": [
        "W = torch.arange(9).view(3, 3) # We can also stack multiple operations in a single line\n",
        "print(\"W\", W)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S7fkb6OamJpv"
      },
      "source": [
        "## TODO: Perform a matrix multiplication between x and W\n",
        "h = ... \n",
        "print(\"h\", h)\n",
        "# Verify the result by calculating it by hand too!"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wSXfdKWTmJpv"
      },
      "source": [
        "#### Indexing\n",
        "\n",
        "We often have the situation where we need to select a part of a tensor. Indexing works just like in numpy, so let's try it:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "n-wYiEP9mJpw"
      },
      "source": [
        "x = torch.arange(12).view(3, 4)\n",
        "print(\"X\", x)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-ydsFgKImJpw"
      },
      "source": [
        "## TODO: Print the second column of the tensor x\n",
        "print(...) "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "a1e7LFD-mJpx"
      },
      "source": [
        "## TODO: Print the first row of the tensor x\n",
        "print(...)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "e5MHj0AAmJpx"
      },
      "source": [
        "## TODO: Print the elements of the tensor x that are both in first two rows and the last column \n",
        "print(...)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cKX077EcmJpx"
      },
      "source": [
        "## TODO: Print the middle two rows of tensor x\n",
        "print(...)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XafS-KDNmJpy"
      },
      "source": [
        "### Dynamic Computation Graph and Backpropagation\n",
        "\n",
        "One of the main reasons for using PyTorch in Deep Learning projects is that we can automatically get **gradients/derivatives** of functions that we define. We will mainly use PyTorch for implementing neural networks, and they are just fancy functions. If we use weight matrices in our function that we want to learn, then those are called the **parameters** or simply the **weights**.\n",
        "\n",
        "If our neural network would output a single scalar value, we would talk about taking the **derivative**, but you will see that quite often we will have **multiple** output variables (\"values\"); in that case we talk about **gradients**. It's a more general term.\n",
        "\n",
        "Given an input $\\mathbf{x}$, we define our function by **manipulating** that input, usually by matrix-multiplications with weight matrices and additions with so-called bias vectors. As we manipulate our input, we are automatically creating a **computational graph**. This graph shows how to arrive at our output from our input. \n",
        "PyTorch is a **define-by-run** framework; this means that we can just do our manipulations, and PyTorch will keep track of that graph for us. Thus, we create a dynamic computation graph along the way.\n",
        "\n",
        "So, to recap: the only thing we have to do is to compute the **output**, and then we can ask PyTorch to automatically get the **gradients**. \n",
        "\n",
        "> **Note:  Why do we want gradients?** Consider that we have defined a function, a neural net, that is supposed to compute a certain output $y$ for an input vector $\\mathbf{x}$. We then define an **error measure** that tells us how wrong our network is; how bad it is in predicting output $y$ from input $\\mathbf{x}$. Based on this error measure, we can use the gradients to **update** the weights $\\mathbf{W}$ that were responsible for the output, so that the next time we present input $\\mathbf{x}$ to our network, the output will be closer to what we want.\n",
        "\n",
        "The first thing we have to do is to specify which tensors require gradients. By default, when we create a tensor, it does not require gradients."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Facxlx-umJpy"
      },
      "source": [
        "x = torch.ones((3,))\n",
        "print(x.requires_grad)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0-8dR8A5mJpy"
      },
      "source": [
        "We can change this for an existing tensor using the function `requires_grad_()` (underscore indicating that this is a in-place operation). Alternatively, when creating a tensor, you can pass the argument `requires_grad=True` to most initializers we have seen above."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gssDovRXmJpy"
      },
      "source": [
        "x.requires_grad_(True)\n",
        "print(x.requires_grad)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GZ45PmZ2mJpz"
      },
      "source": [
        "In order to get familiar with the concept of a computation graph, we will create one for the following function:\n",
        "\n",
        "$$y = \\frac{1}{|x|}\\sum_i \\left[(x_i + 4)^2 + 9\\right]$$\n",
        "\n",
        "You could imagine that $x$ are our parameters, and we want to optimize (either maximize or minimize) the output $y$. For this, we want to obtain the gradients $\\partial y / \\partial \\mathbf{x}$. For our example, we'll use $\\mathbf{x}=[0,1,2]$ as our input."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1LpQn_HZmJpz"
      },
      "source": [
        "## TODO: Create a tensor with values [0.0, 1.0, 2.0] which is requires gradients\n",
        "# Note that only float tensors can have gradients\n",
        "x = ...\n",
        "print(\"X\", x)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wRhZ4tAUmJpz"
      },
      "source": [
        "Now let's build the computation graph step by step. You can combine multiple operations in a single line, but we will separate them here to get a better understanding of how each operation is added to the computation graph."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jpxordqtmJpz"
      },
      "source": [
        "## TODO: Implement the function above in four steps, each representing one single operation (add -> a, square -> b, add -> c, mean -> y)\n",
        "a = ...\n",
        "b = ...\n",
        "c = ...\n",
        "y = ...\n",
        "print(\"Y\", y)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZDKsddoUmJpz"
      },
      "source": [
        "Using the statements above, we have created a computation graph that looks similar to the figure below:\n",
        "\n",
        "<center style=\"width: 100%\"><img src=\"https://github.com/adityaintwala/ACM-GeometricDeepLearning-2021/blob/b10d477bc9ffa98d399e7ab38e3fa496b052a707/Data/Doc/dynamic_computation_graph.PNG?raw=1\" width=\"200px\"></center>\n",
        "\n",
        "We calculate $a$ based on the inputs $x$ and the constant $4$, $b$ is $a$ squared, and so on. The visualization is an abstraction of the dependencies between inputs and outputs of the operations we have applied.\n",
        "Each node of the computation graph has automatically defined a function for calculating the gradients with respect to its inputs, `grad_fn`. You can see this when we printed the output tensor $y$. This is why the computation graph is usually visualized in the reverse direction (arrows point from the result to the inputs). We can perform backpropagation on the computation graph by calling the function `backward()` on the last output, which effectively calculates the gradients for each tensor that has the property `requires_grad=True`:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xaCPYOtLmJp0"
      },
      "source": [
        "y.backward()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qcHPYQXRmJp0"
      },
      "source": [
        "`x.grad` will now contain the gradient $\\partial y/ \\partial \\mathcal{x}$, and this gradient indicates how a change in $\\mathbf{x}$ will affect output $y$ given the current input $\\mathbf{x}=[0,1,2]$:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v8DdHClOmJp0"
      },
      "source": [
        "print(x.grad)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IVlGuT06mJp0"
      },
      "source": [
        "We can also verify these gradients by hand. We will calculate the gradients using the chain rule, in the same way as PyTorch did it:\n",
        "\n",
        "$$\\frac{\\partial y}{\\partial x_i} = \\frac{\\partial y}{\\partial c_i}\\frac{\\partial c_i}{\\partial b_i}\\frac{\\partial b_i}{\\partial a_i}\\frac{\\partial a_i}{\\partial x_i}$$\n",
        "\n",
        "Note that we have simplified this equation to index notation, and by using the fact that all operation besides the mean do not combine the elements in the tensor. The partial derivatives are:\n",
        "\n",
        "$$\n",
        "\\frac{\\partial a_i}{\\partial x_i} = 1,\\hspace{1cm}\n",
        "\\frac{\\partial b_i}{\\partial a_i} = 2\\cdot a_i\\hspace{1cm}\n",
        "\\frac{\\partial c_i}{\\partial b_i} = 1\\hspace{1cm}\n",
        "\\frac{\\partial y}{\\partial c_i} = \\frac{1}{3}\n",
        "$$\n",
        "\n",
        "Hence, with the input being $\\mathbf{x}=[0,1,2]$, our gradients are $\\partial y/\\partial \\mathbf{x}=[8/3,10/3,4]$. The previous code cell should have printed the same result."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "epWXbjpYmJp0"
      },
      "source": [
        "### GPU support\n",
        "\n",
        "A crucial feature of PyTorch is the support of GPUs, short for Graphics Processing Unit. A GPU can perform many thousands of small operations in parallel, making it very well suitable for performing large matrix operations in neural networks. When comparing GPUs to CPUs, we can list the following main differences (credit: [Kevin Krewell, 2009](https://blogs.nvidia.com/blog/2009/12/16/whats-the-difference-between-a-cpu-and-a-gpu/)) \n",
        "\n",
        "<center style=\"width: 100%\"><img src=\"https://github.com/adityaintwala/ACM-GeometricDeepLearning-2021/blob/main/Data/Doc/CPU_GPU.PNG?raw=1\" width=\"700px\"></center>\n",
        "\n",
        "CPUs and GPUs have both different advantages and disadvantages, which is why many computers contain both components and use them for different tasks. In case you are not familiar with GPUs, you can read up more details in this [NVIDIA blog post](https://blogs.nvidia.com/blog/2009/12/16/whats-the-difference-between-a-cpu-and-a-gpu/) or [here](https://www.intel.com/content/www/us/en/products/docs/processors/what-is-a-gpu.html). \n",
        "\n",
        "GPUs can accelerate the training of your network up to a factor of $100$ which is essential for large neural networks. PyTorch implements a lot of functionality for supporting GPUs (mostly those of NVIDIA due to the libraries [CUDA](https://developer.nvidia.com/cuda-zone) and [cuDNN](https://developer.nvidia.com/cudnn)). First, let's check whether you have a GPU available:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0fTXGN5ZmJp1"
      },
      "source": [
        "gpu_avail = torch.cuda.is_available()\n",
        "print(\"Is the GPU available? %s\" % str(gpu_avail))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j2VQO5PamJp1"
      },
      "source": [
        "If you have a GPU on your computer but the command above returns False, make sure you have the correct CUDA-version installed. The `dl2020` environment comes with the CUDA-toolkit 10.1, which is selected for the Lisa supercomputer. Please change it if necessary (CUDA 10.2 is currently common). On Google Colab, make sure that you have selected a GPU in your runtime setup (in the menu, check under `Runtime -> Change runtime type`). \n",
        "\n",
        "By default, all tensors you create are stored on the CPU. We can push a tensor to the GPU by using the function `.to(...)`, or `.cuda()`. However, it is often a good practice to define a `device` object in your code which points to the GPU if you have one, and otherwise to the CPU. Then, you can write your code with respect to this device object, and it allows you to run the same code on both a CPU-only system, and one with a GPU. Let's try it below. We can specify the device as follows: "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b7M--pi6mJp1"
      },
      "source": [
        "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
        "print(\"Device\", device)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Il7fja3KmJp1"
      },
      "source": [
        "Now let's create a tensor and push it to the device:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BIkGKHcwmJp1"
      },
      "source": [
        "x = torch.zeros(2, 3)\n",
        "\n",
        "## TODO: Push the tensor to \"device\"\n",
        "x = ...\n",
        "print(\"X\", x)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wgmirjG7mJp2"
      },
      "source": [
        "In case you have a GPU, you should now see the attribute `device='cuda:0'` being printed next to your tensor. The zero next to cuda indicates that this is the zero-th GPU device on your computer. PyTorch also supports multi-GPU systems, but this you will only need once you have very big networks to train (if interested, see the [PyTorch documentation](https://pytorch.org/docs/stable/distributed.html#distributed-basics)). We can also compare the runtime of a large matrix multiplication on the CPU with a operation on the GPU:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O9H3uBJHmJp2"
      },
      "source": [
        "x = torch.randn(5000, 5000)\n",
        "\n",
        "## CPU version\n",
        "start_time = time.time()\n",
        "_ = torch.matmul(x, x)\n",
        "end_time = time.time()\n",
        "print(\"CPU time: %6.5fs\" % (end_time - start_time))\n",
        "\n",
        "## GPU version\n",
        "x = x.to(device)\n",
        "# The first operation on a CUDA device can be slow as it has to establish a CPU-GPU communication first. \n",
        "# Hence, we run an arbitrary command first without timing it for a fair comparison.\n",
        "if torch.cuda.is_available():\n",
        "    _ = torch.matmul(x*0.0, x)\n",
        "start_time = time.time()\n",
        "_ = torch.matmul(x, x)\n",
        "end_time = time.time()\n",
        "print(\"GPU time: %6.5fs\" % (end_time - start_time))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uCoiPRHemJp2"
      },
      "source": [
        "Depending on the size of the operation and the CPU/GPU in your system, the speedup of this operation can be >500x. As `matmul` operations are very common in neural networks, we can already see the great benefit of training a NN on a GPU. The time estimate can be relatively noisy here because we haven't run it for multiple times. Feel free to extend this, but it also takes longer to run.\n",
        "\n",
        "When generating random numbers, the seed between CPU and GPU is not synchronized. Hence, we need to set the seed on the GPU separately to ensure a reproducible code. Note that due to different GPU architectures, running the same code on different GPUs does not guarantee the same random numbers. Still, we don't want that our code gives us a different output every time we run it on the exact same hardware. Hence, we also set the seed on the GPU:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4auaYoIFmJp2"
      },
      "source": [
        "# GPU operations have a separate seed we also want to set\n",
        "if torch.cuda.is_available(): \n",
        "    torch.cuda.manual_seed(7)\n",
        "    torch.cuda.manual_seed_all(7)\n",
        "    \n",
        "# Additionally, some operations on a GPU are implemented stochastic for efficiency\n",
        "# We want to ensure that all operations are deterministic on GPU (if used) for reproducibility\n",
        "torch.backends.cudnn.determinstic = True\n",
        "torch.backends.cudnn.benchmark = False"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RsKvc9ptmJp2"
      },
      "source": [
        "## Learning by example: Continuous XOR\n",
        "\n",
        "If we want to build a neural network in PyTorch, we could specify all our parameters (weight matrices, bias vectors) using `Tensors` (with `requires_grad=True`), ask PyTorch to calculate the gradients and then adjust the parameters. But things can quickly get cumbersome if we have a lot of parameters. In PyTorch, there is a package called `torch.nn` that makes building neural networks more convenient. \n",
        "\n",
        "We will introduce the libraries and all additional parts you might need to train a neural network in PyTorch, using a simple example classifier on a simple yet well known example: XOR. Given two binary inputs $x_1$ and $x_2$, the label to predict is $1$ if either $x_1$ or $x_2$ is $1$ while the other is $0$, or the label is $0$ in all other cases. The example became famous by the fact that a single neuron, i.e. a linear classifier, cannot learn this simple function.\n",
        "Hence, we will learn how to build a small neural network that can learn this function. \n",
        "To make it a little bit more interesting, we move the XOR into continuous space and introduce some gaussian noise on the binary inputs. Our desired separation of an XOR dataset could look as follows:\n",
        "\n",
        "<center style=\"width: 100%\"><img src=\"https://github.com/adityaintwala/ACM-GeometricDeepLearning-2021/blob/main/Data/Doc/Xor.PNG?raw=1\" width=\"350px\"></center>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hJlASM8fmJp3"
      },
      "source": [
        "### The model\n",
        "\n",
        "The package `torch.nn` defines a series of useful classes like linear networks layers, activation functions, loss functions etc. A full list can be found [here](https://pytorch.org/docs/stable/nn.html). In case you need a certain network layer, check the documentation of the package first before writing the layer yourself as the package likely contains the code for it already. We import it below:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5TRhAPAtmJp3"
      },
      "source": [
        "import torch.nn as nn"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MoiLa_Q_mJp3"
      },
      "source": [
        "Additionally to `torch.nn`, there is also `torch.nn.functional`. It contains functions that are used in network layers. This is in contrast to `torch.nn` which defines them as `nn.Modules` (more on it below), and `torch.nn` actually uses a lot of functionalities from `torch.nn.functional`. Hence, the functional package is useful in many situations, and so we import it as well here."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9j7UIsUhmJp4"
      },
      "source": [
        "import torch.nn.functional as F"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L1jb4kYzmJp4"
      },
      "source": [
        "#### nn.Module\n",
        "\n",
        "In PyTorch, a neural network is build up out of modules. Modules can contain other modules, and a neural network is considered to be a module itself as well. The basic template of a module is as follows:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EQHwORADmJp5"
      },
      "source": [
        "class MyModule(nn.Module):\n",
        "    \n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        # Some init for my module\n",
        "        \n",
        "    def forward(self, x):\n",
        "        # Function for performing the calculation of the module.\n",
        "        pass"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eZng3I-1mJp5"
      },
      "source": [
        "The forward function is where the computation of the module is taken place, and is executed when you call the module (`nn = MyModule(); nn(x)`). In the init function, we usually create the parameters of the module, using `nn.Parameter`, or defining other modules that are used in the forward function. The backward calculation is done automatically, but could be overwritten as well if wanted.\n",
        "\n",
        "#### Simple classifier\n",
        "We can now make use of the pre-defined modules in the `torch.nn` package, and define our own small neural network. We will use a minimal network with a input layer, one hidden layer with tanh as activation function, and a output layer. In other words, our networks should look something like this:\n",
        "\n",
        "<center width=\"100%\"><img src=\"https://github.com/adityaintwala/ACM-GeometricDeepLearning-2021/blob/main/Data/Doc/NeuralNetwork.PNG?raw=1\" width=\"300px\"></center>\n",
        "\n",
        "The input neurons are shown in blue, which represent the coordinates $x_1$ and $x_2$ of a data point. The hidden neurons including a tanh activation are shown in white, and the output neuron in red.\n",
        "In PyTorch, we can define this as follows:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rlXqEd_amJp6"
      },
      "source": [
        "class SimpleClassifier(nn.Module):\n",
        "    \n",
        "    def __init__(self, num_inputs, num_hidden, num_outputs):\n",
        "        super().__init__()\n",
        "        # Initialize the modules we need to build the network\n",
        "        ## TODO: Implement the linear layers using nn.Linear(input_neurons, output_neurons)\n",
        "        self.linear1 = ...\n",
        "        self.linear2 = ...\n",
        "        ## TODO: Implement the activation function using nn.Tanh\n",
        "        self.act_fn = ...\n",
        "        \n",
        "        \n",
        "    def forward(self, x):\n",
        "        # Perform the calculation of the model to determine the prediction\n",
        "        ## TODO: Execute the NN modules in sequence to build the neural network computation graph\n",
        "        x = ...\n",
        "        return x"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0DaUvxIxmJp6"
      },
      "source": [
        "For the examples in this notebook, we will use a tiny neural network with two input neurons and four hidden neurons. As we perform binary classification, we will use a single output neuron. Note that we do not apply a sigmoid on the output yet. This is because other functions, especially the loss, are more efficient and precise to calculate on the original outputs instead of the sigmoid output. We will discuss the detailed reason later."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "y7DpmTcNmJp6"
      },
      "source": [
        "model = SimpleClassifier(num_inputs=2, num_hidden=4, num_outputs=1)\n",
        "# Printing a module shows all its submodules\n",
        "print(model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qnHuzgT5mJp6"
      },
      "source": [
        "Printing the model lists all submodules it contains. The parameters of a module can be obtained by using its `parameters()` functions, or `named_parameters()` to get a name to each parameter object. For our small neural network, we have the following parameters:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2NiZMYc2mJp7"
      },
      "source": [
        "for name, param in model.named_parameters():\n",
        "    print(\"Parameter %s, shape %s\" % (name, str(param.shape)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4_NQAya9mJp7"
      },
      "source": [
        "Each linear layer has a weight matrix of the shape `[output, input]`, and a bias of the shape `[output]`. The tanh activation function does not have any parameters. Note that parameters are only registered for `nn.Module` objects that are direct object attributes, i.e. `self.a = ...`. If you define a list of modules, the parameters of those are not registered for the outer module and can cause some issues when you try to optimize your module. There are alternatives, like `nn.ModuleList`, `nn.ModuleDict` and `nn.Sequential`, that allow you to have different data structures of modules. We will use them in a few later tutorials and explain them there. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xPfGyYUtmJp8"
      },
      "source": [
        "### The data\n",
        "\n",
        "PyTorch also provides a few functionalities to load the training and test data efficiently, summarized in the package `torch.utils.data`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zsM6W3PemJp8"
      },
      "source": [
        "import torch.utils.data as data"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uxN1zXsjmJp9"
      },
      "source": [
        "The data package defines two classes which are the standard interface for handling data in PyTorch: `data.Dataset`, and `data.DataLoader`. The dataset class provides an uniform interface to access the training/test data, while the data loader makes sure to efficiently load and stack the data points from the dataset into batches during training."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mjS31RRgmJp9"
      },
      "source": [
        "#### The dataset class\n",
        "\n",
        "The dataset class summarizes the basic functionality of a dataset in a natural way. To define a dataset in PyTorch, we simply specify two functions: `__getitem__`, and `__len__`. The get-item function has to return the $i$-th data point in the dataset, while the len function returns the size of the dataset. For the XOR dataset, we can define the dataset class as follows:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C83-he_vmJp9"
      },
      "source": [
        "class XORDataset(data.Dataset):\n",
        "    \n",
        "    def __init__(self, size, std=0.1):\n",
        "        \"\"\"\n",
        "        Inputs:\n",
        "            size - Number of data points we want to generate\n",
        "            std - Standard deviation of the noise (see generate_continuous_xor function)\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        self.size = size\n",
        "        self.std = std\n",
        "        self.generate_continuous_xor()\n",
        "        \n",
        "    \n",
        "    def generate_continuous_xor(self):\n",
        "        # Each data point in the XOR dataset has two variables, x and y, that can be either 0 or 1\n",
        "        # The label is their XOR combination, i.e. 1 if only x or only y is 1 while the other is 0.\n",
        "        # If x=y, the label is 0.\n",
        "        data = torch.randint(low=0, high=2, size=(self.size, 2), dtype=torch.float32)\n",
        "        label = (data.sum(dim=1) == 1).to(torch.long)\n",
        "        # To make it slightly more challenging, we add a bit of gaussian noise to the data points.\n",
        "        data += self.std * torch.randn(data.shape)\n",
        "        \n",
        "        self.data = data\n",
        "        self.label = label\n",
        "        \n",
        "    \n",
        "    def __len__(self):\n",
        "        ## TODO: Return the number of data points the data has\n",
        "        return ...\n",
        "    \n",
        "    \n",
        "    def __getitem__(self, idx):\n",
        "        ## TODO: Return the idx-th data point and label of the dataset\n",
        "        data_point = ...\n",
        "        data_label = ...\n",
        "        return data_point, data_label"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UCMfkOLwmJp-"
      },
      "source": [
        "Let's try to create such a dataset and inspect it:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FYVYtTwdmJp-"
      },
      "source": [
        "dataset = XORDataset(size=200)\n",
        "print(\"Size of dataset:\", len(dataset))\n",
        "print(\"Data point 0:\", dataset[0])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QFylecSymJp-"
      },
      "source": [
        "To better relate to the dataset, we visualize the samples below. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QViF4Jz0mJp_"
      },
      "source": [
        "def visualize_samples(data, label):\n",
        "    if isinstance(data, torch.Tensor):\n",
        "        data = data.cpu().numpy()\n",
        "    if isinstance(label, torch.Tensor):\n",
        "        label = label.cpu().numpy()\n",
        "    data_0 = data[label == 0]\n",
        "    data_1 = data[label == 1]\n",
        "    \n",
        "    plt.figure(figsize=(4,4))\n",
        "    plt.scatter(data_0[:,0], data_0[:,1], edgecolor=\"#333\", label=\"Class 0\")\n",
        "    plt.scatter(data_1[:,0], data_1[:,1], edgecolor=\"#333\", label=\"Class 1\")\n",
        "    plt.title(\"Dataset samples\")\n",
        "    plt.ylabel(r\"$x_2$\")\n",
        "    plt.xlabel(r\"$x_1$\")\n",
        "    plt.legend()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c1RiuvqcmJp_"
      },
      "source": [
        "visualize_samples(dataset.data, dataset.label)\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3oWriws6mJp_"
      },
      "source": [
        "#### The data loader class\n",
        "\n",
        "The class `torch.utils.data.DataLoader` represents a Python iterable over a dataset with support for automatic batching, multi-process data loading and many more features. The data loader communicates with the dataset using the function `__getitem__`, and stacks its outputs as tensors over the first dimension to form a batch.\n",
        "In contrast to the dataset class, we usually don't have to define our own data loader class, but can just create an object of it with the dataset as input. Additionally, we can configure our data loader with the following input arguments (only a selection, see full list [here](https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader)):\n",
        "\n",
        "* `batch_size`: Number of samples to stack per batch\n",
        "* `shuffle`: If True, the data is returned in a random order. This is important during training for introducing stochasticity. \n",
        "* `num_workers`: Number of subprocesses to use for data loading. The default, 0, means that the data will be loaded in the main process which can slow down training for datasets where loading a data point takes a considerable amount of time (e.g. large images). More workers are recommended for those, but can cause issues on Windows computers. For tiny datasets as ours, 0 workers are usually faster.\n",
        "* `pin_memory`: If True, the data loader will copy Tensors into CUDA pinned memory before returning them. This can save some time for large data points on GPUs. Usually a good practice to use for a training set, but not necessarily for validation and test to save memory on the GPU.\n",
        "* `drop_last`: If True, the last batch is dropped in case it is smaller than the specified batch size. This occurs when the dataset size is not a multiple of the batch size. Only potentially helpful during training to keep a consistent batch size.\n",
        "\n",
        "Let's create a simple data loader below:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6l0Eihh6mJqA"
      },
      "source": [
        "data_loader = data.DataLoader(dataset, batch_size=8, shuffle=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "L3ls8hT5mJqA"
      },
      "source": [
        "# next(iter(...)) catches the first batch of the data loader\n",
        "# If shuffle is True, this will return a different batch every time we run this cell\n",
        "# For iterating over the whole dataset, we can simple use \"for batch in data_loader: ...\"\n",
        "data_inputs, data_labels = next(iter(data_loader))\n",
        "\n",
        "# The shape of the outputs are [batch_size, d_1,...,d_N] where d_1,...,d_N are the \n",
        "# dimensions of the data point returned from the dataset class\n",
        "print(\"Data inputs\", data_inputs.shape, \"\\n\", data_inputs)\n",
        "print(\"Data labels\", data_labels.shape, \"\\n\", data_labels)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ggwnur38mJqA"
      },
      "source": [
        "### Optimization\n",
        "\n",
        "After defining the model and the dataset, it is time to prepare the optimization of the model. During training, we will perform the following steps:\n",
        "\n",
        "1. Get a batch from the data loader\n",
        "2. Obtain the predictions from the model for the batch\n",
        "3. Calculate the loss based on the difference between predictions and labels\n",
        "4. Backpropagation: calculate the gradients for every parameter with respect to the loss\n",
        "5. Update the parameters of the model in the direction of the gradients\n",
        "\n",
        "We have seen how we can do step 1, 2 and 4 in PyTorch. Now, we will look at step 3 and 5."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sfIXs_b5mJqB"
      },
      "source": [
        "#### Loss modules\n",
        "\n",
        "We can calculate the loss for a batch by simply performing a few tensor operations as those are automatically added to the computation graph. For instance, for binary classification, we can use Binary Cross Entropy (BCE) which is defined as follows:\n",
        "\n",
        "$$\\mathcal{L}_{BCE} = \\sum_i y_i \\log x_i + (1 - y_i) \\log (1 - x_i)$$\n",
        "\n",
        "where $y$ are our labels, and $x$ our predictions, both in the range of $[0,1]$. However, PyTorch already provides a list of predefined loss functions which we can use (see [here](https://pytorch.org/docs/stable/nn.html#loss-functions) for a full list). For instance, for BCE, PyTorch has two modules: `nn.BCELoss()`, `nn.BCEWithLogitsLoss()`. While `nn.BCELoss` expects the inputs $x$ to be in the range $[0,1]$, i.e. the output of a sigmoid, `nn.BCEWithLogitsLoss` combines a sigmoid layer and the BCE loss in a single class. This version is numerically more stable than using a plain Sigmoid followed by a BCE loss because of the logarithms applied in the loss function. Hence, it is adviced to use loss functions applied on \"logits\" where possible (remember to not apply a sigmoid on the output of the model in this case!). For our model defined above, we therefore use the module `nn.BCEWithLogitsLoss`. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "l7o01SmOmJqB"
      },
      "source": [
        "loss_module = nn.BCEWithLogitsLoss()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RtixEHqqmJqB"
      },
      "source": [
        "#### Stochastic Gradient Descent\n",
        "\n",
        "For updating the parameters, PyTorch provides the package `torch.optim` that has most popular optimizers implemented. We will discuss the specific optimizers and their differences later in the course, but will for now use the simplest of them: `torch.optim.SGD`. Stochastic Gradient Descent updates parameters by multiplying the gradients with a small constant, called learning rate, and subtracting those from the parameters (hence minimizing the loss). Therefore, we slowly move towards the direction of minimizing the loss. A good default value of the learning rate for a small network as ours is 0.1. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "InIU_8XMmJqC"
      },
      "source": [
        "# Input to the optimizer are the parameters of the model: model.parameters()\n",
        "optimizer = torch.optim.SGD(model.parameters(), lr=0.1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NXdGUPfYmJqC"
      },
      "source": [
        "The optimizer provides two useful functions: `optimizer.step()`, and `optimizer.zero_grad()`. The step function updates the parameters based on the gradients as explained above. The function `optimizer.zero_grad()` sets the gradients of all parameters to zero. While this function seems less relevant at first, it is a crucial pre-step before performing backpropagation. If we would call the `backward` function on the loss while the parameter gradients are non-zero from the previous batch, the new gradients would actually be added to the previous ones instead of overwriting them. This is done because a parameter might occur multiple times in a computation graph, and we need to sum the gradients in this case instead of replacing them. Hence, remember to call `optimizer.zero_grad()` before calculating the gradients of a batch."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dtqyEhuUmJqC"
      },
      "source": [
        "### Training\n",
        "\n",
        "Finally, we are ready to train our model. As a first step, we create a slightly larger dataset and specify a data loader with a larger batch size. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VoxAXngYmJqC"
      },
      "source": [
        "train_dataset = XORDataset(size=1000)\n",
        "train_data_loader = data.DataLoader(train_dataset, batch_size=128, shuffle=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yDewhVQWmJqD"
      },
      "source": [
        "Now, we can write a small training function. Remember our five steps: load a batch, obtain the predictions, calculate the loss, backpropagate, and update. Additionally, we have to push all data and model parameters to the device of our choice (GPU if available). For the tiny neural network we have, communicating the data to the GPU actually takes much more time than we could save from running the operation on GPU. For large networks, the communication time is significantly smaller than the actual runtime making a GPU crucial in these cases. Still, to practice, we will push the data to GPU here. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "i1Ksclt_mJqD"
      },
      "source": [
        "# Push model to device. Has to be only done once\n",
        "model.to(device)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O5bhX5ffmJqD"
      },
      "source": [
        "In addition, we set our model to training mode. This is done by calling `model.train()`. There exist certain modules that need to perform a different forward step during training than during testing (e.g. BatchNorm and Dropout), and we can switch between them using `model.train()` and `model.eval()`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "e24Hg_ApmJqD"
      },
      "source": [
        "def train_model(model, optimizer, data_loader, loss_module, num_epochs=100):\n",
        "    # Set model to train mode\n",
        "    model.train() \n",
        "    \n",
        "    # Training loop\n",
        "    for epoch in range(num_epochs):\n",
        "        print('============= epoch {} ============='.format(epoch))\n",
        "        for data_inputs, data_labels in data_loader:\n",
        "            \n",
        "            ## TODO: Step 1 - Move input data to device\n",
        "            data_inputs, data_labels = ...\n",
        "            \n",
        "            ## TODO: Step 2 - Run the model on the input data\n",
        "            preds = ...\n",
        "            \n",
        "            ## TODO: Step 3 - Calculate the loss using the module loss_module\n",
        "            loss = ...\n",
        "            print('============= epoch {} : loss {} ============='.format(epoch, loss))\n",
        "            ## TODO: Step 4 - Perform backpropagation\n",
        "            ...\n",
        "            \n",
        "            ## TODO: Step 5 - Update parameters\n",
        "            ..."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cbl7gXmlmJqD"
      },
      "source": [
        "train_model(model, optimizer, train_data_loader, loss_module)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Sh-F2kS6mJqE"
      },
      "source": [
        "#### Saving a model\n",
        "\n",
        "After finish training a model, we save the model to disk so that we can load the same weights at a later time. For this, we extract the so-called `state_dict` from the model which contains all learnable parameters. For our simple model, the state dict contains the following entries:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "k9RxJ53WmJqE"
      },
      "source": [
        "state_dict = model.state_dict()\n",
        "print(state_dict)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GjDUyrZtmJqE"
      },
      "source": [
        "To save the state dictionary, we can use `torch.save`:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3AhSpGhPmJqE"
      },
      "source": [
        "# torch.save(object, filename). For the filename, any extension can be used\n",
        "torch.save(state_dict, \"our_model.tar\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X0UtccmFmJqE"
      },
      "source": [
        "To load a model from a state dict, we use the function `torch.load` to load the state dict from the disk, and the module function `load_state_dict` to overwrite our parameters with the new values:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TloDFJiJmJqE"
      },
      "source": [
        "# Load state dict from the disk (make sure it is the same name as above)\n",
        "state_dict = torch.load(\"our_model.tar\")\n",
        "\n",
        "# Create a new model and load the state\n",
        "new_model = SimpleClassifier(num_inputs=2, num_hidden=4, num_outputs=1)\n",
        "new_model.load_state_dict(state_dict)\n",
        "\n",
        "# Verify that the parameters are the same\n",
        "print(\"Original model\\n\", model.state_dict())\n",
        "print(\"\\nLoaded model\\n\", new_model.state_dict())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CupSbGPvmJqF"
      },
      "source": [
        "A detailed tutorial on saving and loading models in PyTorch can be found [here](https://pytorch.org/tutorials/beginner/saving_loading_models.html)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "osOyNndimJqF"
      },
      "source": [
        "### Evaluation\n",
        "\n",
        "Once we have trained a model, it is time to evaluate it on a held-out test set. As our dataset consist of randomly generated data points, we need to first create a test set with a corresponding data loader."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7M3v4oDtmJqF"
      },
      "source": [
        "test_dataset = XORDataset(size=500)\n",
        "# drop_last -> Don't drop the last batch although it is smaller than 128\n",
        "test_data_loader = data.DataLoader(test_dataset, batch_size=128, shuffle=False, drop_last=False) "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I-nBwNpGmJqF"
      },
      "source": [
        "As metric, we will use accuracy which is calculated as follows:\n",
        "\n",
        "$$acc = \\frac{\\#\\text{correct predictions}}{\\#\\text{all predictions}} = \\frac{TP+TN}{TP+TN+FP+FN}$$\n",
        "\n",
        "where TP are the true positives, TN true negatives, FP false positives, and FN the fale negatives. \n",
        "\n",
        "When evaluating the model, we don't need to keep track of the computation graph as we don't intend to calculate the gradients. This reduces the required memory and speed up the model. In PyTorch, we can deactivate the computation graph using `with torch.no_grad(): ...`. Remember to additionally set the model to eval mode."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EiTv1q-nmJqF"
      },
      "source": [
        "def eval_model(model, data_loader):\n",
        "    model.eval() # Set model to eval mode\n",
        "    true_preds, num_preds = 0., 0.\n",
        "    \n",
        "    with torch.no_grad(): # Deactivate gradients for the following code\n",
        "        for data_inputs, data_labels in data_loader:\n",
        "            \n",
        "            ## TODO: Step 1 - Move data to device\n",
        "            data_inputs, data_labels = data_inputs.to(device), data_labels.to(device)\n",
        "            \n",
        "            ## TODO: Step 2 - Run the model on the input data\n",
        "            preds = ...\n",
        "            \n",
        "            ## TODO: Step 3 - Determine binary predictions of the model\n",
        "            pred_labels = ...\n",
        "            \n",
        "            ## TODO: Step 4 - Keep records of predictions for the accuracy metric (true_preds=TP+TN, num_preds=TP+TN+FP+FN)\n",
        "            true_preds += ...\n",
        "            num_preds += ...\n",
        "            \n",
        "    acc = true_preds / num_preds\n",
        "    print(\"Accuracy of the model: %4.2f%%\" % (100.0*acc))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xef-U9pWmJqG"
      },
      "source": [
        "eval_model(model, test_data_loader)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "spzqqn-gmJqG"
      },
      "source": [
        "If we trained our model correctly, we should see a score close to 100% accuracy. However, this is only possible because of our simple task, and unfortunately, we usually don't get such high scores on test sets of more complex tasks."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sTAqvA0QmJqG"
      },
      "source": [
        "#### Visualizing classification boundaries\n",
        "\n",
        "To visualize what our model has learned, we can perform a prediction for every data point in a range of $[-0.5, 1.5]$, and visualize the predicted class as in the sample figure at the beginning of this section. This shows where the model has created decision boundaries, and which points would be classified as $0$, and which as $1$. We therefore get a background image out of blue (class 0) and orange (class 1). The spots where the model is uncertain we will see a blurry overlap. The specific code is less relevant compared to the output figure which should hopefully show us a clear separation of classes:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rXuMK9CxmJqG"
      },
      "source": [
        "@torch.no_grad() # Decorator, same effect as \"with torch.no_grad(): ...\" over the whole function.\n",
        "def visualize_classification(model, data, label):\n",
        "    if isinstance(data, torch.Tensor):\n",
        "        data = data.cpu().numpy()\n",
        "    if isinstance(label, torch.Tensor):\n",
        "        label = label.cpu().numpy()\n",
        "    data_0 = data[label == 0]\n",
        "    data_1 = data[label == 1]\n",
        "    \n",
        "    plt.figure(figsize=(4,4))\n",
        "    plt.scatter(data_0[:,0], data_0[:,1], edgecolor=\"#333\", label=\"Class 0\")\n",
        "    plt.scatter(data_1[:,0], data_1[:,1], edgecolor=\"#333\", label=\"Class 1\")\n",
        "    plt.title(\"Dataset samples\")\n",
        "    plt.ylabel(r\"$x_2$\")\n",
        "    plt.xlabel(r\"$x_1$\")\n",
        "    plt.legend()\n",
        "    \n",
        "    # Let's make use of a lot of operations we have learned above\n",
        "    model.to(device)\n",
        "    c0 = torch.Tensor(to_rgba(\"C0\")).to(device)\n",
        "    c1 = torch.Tensor(to_rgba(\"C1\")).to(device)\n",
        "    x1 = torch.arange(-0.5, 1.5, step=0.01, device=device)\n",
        "    x2 = torch.arange(-0.5, 1.5, step=0.01, device=device)\n",
        "    xx1, xx2 = torch.meshgrid(x1, x2) # Meshgrid function as in numpy\n",
        "    model_inputs = torch.stack([xx1, xx2], dim=-1)\n",
        "    preds = model(model_inputs)\n",
        "    preds = torch.sigmoid(preds)\n",
        "    output_image = preds * c0[None,None] + (1 - preds) * c1[None,None] # Specifying \"None\" in a dimension creates a new one\n",
        "    output_image = output_image.cpu().numpy() # Convert to numpy array. This only works for tensors on CPU, hence first push to CPU\n",
        "    plt.imshow(output_image, origin='upper', extent=(-0.5, 1.5, -0.5, 1.5))\n",
        "    plt.grid(False)\n",
        "\n",
        "visualize_classification(model, dataset.data, dataset.label)\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-cyPSTdumJqH"
      },
      "source": [
        "The decision boundaries might not look exactly as in the figure in the preamble of this section which can be caused by running it on CPU or a different GPU architecture. Nevertheless, the result on the accuracy metric should be the approximately the same. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YV7WW_PrmJqH"
      },
      "source": [
        "## Pytorch Tutorials\n",
        "\n",
        "we recommend the official [tutorial website](https://pytorch.org/tutorials/) that contains many tutorials on various topics."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IkqZfonRhoSj"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}