PyTorchTutorial.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8qwzo75IiSdF"
      },
      "source": [
        "# **Don't panic!** A gentle introduction to PyTorch 🔥\n",
        "\n",
        "**Author**: Mikel Malagón <mikel.malagon@ehu.eus>\n",
        "\n",
        "This is an introductory notebook to get up and running with [PyTorch](https://pytorch.org/), one of the most popular deep learning libraries. The aim of this notebook is to introduce the most basic and commonly used aspects of PyTorch to start developing your projects using this library as fast as possible. The notebook assumes basic knowledge on python, machine learning, and that PyTorch is installed, for installation instructions visit the \"install pytorch\" section of its [main page](https://pytorch.org/).\n",
        "\n",
        "⚠️ In case you are running this notebook on Google Colab, I highly recommend enabling GPU in: Runtime > Select runtime type > Hardware accelerator > GPU.\n",
        "\n",
        "<br>\n",
        "\n",
        "**NOTE:** In case you found some bug or have any question/suggestion, don't hesitate to send a mail to the address above.\n",
        "\n",
        "**NOTE:** This notebook is released under the [GPLv3](https://www.gnu.org/licenses/gpl-3.0.en.html) free software license.\n",
        "\n",
        "Thanks to Unai Garciarena for contributing :)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DAzFRJCaoXPW"
      },
      "source": [
        "## Set up\n",
        "\n",
        "First things first, let's import PyTorch to our environment together with matplotlib and numpy."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MTiCoG0dpXq-"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5v9vp0y8oq-f"
      },
      "source": [
        "The following line just checks if there is some GPU available and sets the `device` variable accordingly."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FALQcDznosnm"
      },
      "outputs": [],
      "source": [
        "device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n",
        "print(\"Using device:\", device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "md3MeZ2iim6k"
      },
      "source": [
        "## Introducing the basic element: *the tensor* 🔥💦🌬🗻\n",
        "\n",
        "Tensors are the central object in pytorch. They can have a single dimension (vector), 2 dimensions (matrix), or n-dimensions (tensor). They also share a very similar interface to `ndarray`s of the popular `numpy` library. In PyTorch, tensors are represented with the [`Tensor`](https://pytorch.org/docs/stable/tensors.html) object.\n",
        "\n",
        "Let's create our first pytorch tensors with some default values:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ST8WHqkCiOgt"
      },
      "outputs": [],
      "source": [
        "my_tensor = torch.Tensor([0., 1., 2.]) # Note that 2.0 can be written as 2. in Pyhton\n",
        "my_tensor"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cBIC0eVnvynV"
      },
      "source": [
        "Now that we have created a tensor, let's explore the \"must know\" attributes or properties that all pytorch tensors have:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FtDNycQ8vxz6"
      },
      "outputs": [],
      "source": [
        "print(\"The size of the tensor is:\", my_tensor.size())\n",
        "\n",
        "print(\"The type of the tensor is:\", my_tensor.dtype)\n",
        "\n",
        "print(\"Where is the tensor? \", my_tensor.device)\n",
        "\n",
        "print(\"Has the gradient property? \", my_tensor.requires_grad)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PNEMc7qyr_KM"
      },
      "source": [
        "Note that pytorch tensors have the folowing default properties:\n",
        "- The tensor is stored in CPU memory.\n",
        "- By default numbers are stored as 32 bit float."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XJQN-uTGl3Ml"
      },
      "source": [
        "PyTorch tensors have a very similar interface to numpy's `ndarray`, let's see how to do some basic math operations with tensors:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1JtTeEPwmBvw"
      },
      "outputs": [],
      "source": [
        "double = 2 * my_tensor # multiply a tensor by an scalar\n",
        "half = my_tensor / 2\n",
        "triple = double + my_tensor # operations between tensors\n",
        "\n",
        "print(\"Half:\", half)\n",
        "print(\"Double:\", double)\n",
        "print(\"Triple:\", triple)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_rMHV7FooYvY"
      },
      "source": [
        "Pretty easy! right? As you can see, pytorch tensors can be manipulated using standard python operators, leading to cleaner code. But these operations can also be done directly using pytorch library functions ([reference](https://pytorch.org/docs/stable/torch.html)):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jek8-I_eovpT"
      },
      "outputs": [],
      "source": [
        "# repeat the same operations from the code block above but using pytorch functions\n",
        "double = torch.mul(2., my_tensor)\n",
        "half = torch.div(my_tensor, 2.)\n",
        "triple = torch.add(double, my_tensor)\n",
        "\n",
        "print(\"Half:\", half)\n",
        "print(\"Double:\", double)\n",
        "print(\"Triple:\", triple)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "puaDeZJvpOKK"
      },
      "source": [
        "In the examples above, we have used a tensor created with some default values, but pytorch offers a wide variety of handy constructors to create our tensors. These are some of the most used ones, but the complete list of constructors can be found [here](https://pytorch.org/docs/stable/torch.html#creation-ops).  "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NIQXUXJ2pNSj"
      },
      "outputs": [],
      "source": [
        "# Create a tensor of 3x3 dimension with all values set to zero\n",
        "z = torch.zeros(3, 3)\n",
        "print(\"--> torch.zeros:\\n\\n\", z)\n",
        "\n",
        "# A tensor whose values are sampled from a uniform distribution on the [0, 1) range\n",
        "u = torch.rand((3, 3)) # the shape of the tensor goes between parenthesis here\n",
        "print(\"\\n--> torch.rand:\\n\\n\", u)\n",
        "\n",
        "# A tensor whose values are sampled from a (0, 1) normal\n",
        "n = torch.randn((3, 3))\n",
        "print(\"\\n--> torch.randn:\\n\\n\", n)\n",
        "\n",
        "# Returns a 2-D tensor with ones on the diagonal and zeros elsewhere\n",
        "i = torch.eye(3, 3)\n",
        "print(\"\\n--> torch.eye:\\n\\n\", i)\n",
        "\n",
        "# Returns a 2-D tensor with all ones\n",
        "o = torch.ones(3, 3)\n",
        "print(\"\\n--> torch.ones:\\n\\n\", o)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m1tihpkwrHXC"
      },
      "source": [
        "## Where the witchcraft starts 🧙‍♂️: ✨Autograd✨\n",
        "\n",
        "Remember the code cell in wich we saw the properties of pytorch tensors, the `requires_grad` property of the tensor was set to `False`. This property tells pytorch wheter or not to track the gradient of a tensor. In this section we will see how to use pytorch's Autograd to automatically compute the gradient of a tensor.\n",
        "\n",
        "Let's start by creating a tensor with some default values and setting its `requires_grad` property:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8LoVk7cGrG6m"
      },
      "outputs": [],
      "source": [
        "x = torch.tensor(\n",
        "    [[1., -1.],\n",
        "     [-1., 1.]],\n",
        "    requires_grad=True) # track gradients\n",
        "x"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OAg5yq3QxyLL"
      },
      "source": [
        "Note how this time the `requires_grad` attribute is set to `True`, indicating pytorch to track the operations over it to then be able to compute the gradient. So let's test it! Compute the sum of the squares of the tensor `x`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W7orQhioxt59"
      },
      "outputs": [],
      "source": [
        "x1 = x.pow(2)\n",
        "# print(x1)\n",
        "out = x1.sum()\n",
        "out"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r_QVfMU2yR5D"
      },
      "source": [
        "This is already a bit of a spoiler, notice that the resulting `out` tensor doesn't only contain the final value of the operation (`4.0`) but it also contains a property `grad_fn` that is set to an object named `SumBackward0`. In fact, when operating on a tensor with the `requires_grad` attribute set to `True`, the tensor resulting from any operation on it will contain a `grad_fn` attribute with the derivative function of the applied operation. This way, pytorch is able to construct a graph of the operations that led to a specific tensor to later perform automatic differentiation.\n",
        "\n",
        "In this example, we will compute the gradient of the tensor `x` with respect to the result of the operation, the tensor `out`. By calling the method `backward` on the `out` tensor, Autograd will traverse the described graph and will compute the gradients of all tensors that led to the tensor `out`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jncw5mQmyNRj"
      },
      "outputs": [],
      "source": [
        "out.backward()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h11M5K2J0_kP"
      },
      "source": [
        "To access to the gradients of the tensor `x` we simply have to check its `grad` attribute:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RfproSB61M58"
      },
      "outputs": [],
      "source": [
        "x.grad"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sr1t8_vJ1L93"
      },
      "source": [
        "Voilà! Just using 4 lines of code we used Autograd, the automatic differentiation engine of pytorch, to compute automatically the gradients of a tensor<h1>💃🕺</h1>\n",
        "\n",
        "- 📖 More information on Autograd and how to use it can be found [here](https://pytorch.org/docs/stable/notes/autograd.html)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "thsT8LM3VcWX"
      },
      "source": [
        "TODO (Jon): For those used to TF, it is more natural to use [torch.autograd.grad](https://pytorch.org/docs/stable/generated/torch.autograd.grad.html). I'll include an example comparing that with .backward()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uTZaA81Wy8jv"
      },
      "source": [
        "## Linear regression with gradient descent 📈\n",
        "\n",
        "Now that we have covered the basics of tensors and autograd, we can get into the first proper exercise: linear regression using gradient descent!!!🙌🥳🤓 Not super impressive, but the concepts built in this example perfectly translate to more complex use cases.\n",
        "\n",
        "<br>\n",
        "Let's start by creating some data for the sake of the exercise. Note that the contents of this function aren't very relevant right now, but trying to understand it in depth can be a nice exercise for later (once you have become a pytorch ninja 🥷)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "btW-rXZVz9RE"
      },
      "outputs": [],
      "source": [
        "def data_gen(n_samples=10, x_max=1, intercept=0, noise_gain=0.1):\n",
        "    # uniformly sample the [0, 1) interval\n",
        "    x = x_max*torch.rand(n_samples)\n",
        "    # add some Gaussian noise to x\n",
        "    y = x + noise_gain*torch.randn(n_samples) + intercept\n",
        "    return x[:, None], y[:, None]\n",
        "\n",
        "x, y = data_gen(n_samples=100)\n",
        "plt.scatter(x, y)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GG7RGD4L4mur"
      },
      "source": [
        "Let's create a loss function to measure the error of our regressor. In this case, we'll use the MSE loss, but feel free to change it. Note the similarity of the function to the example used to illustrate Autograd."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZAjjW_Uuy7uv"
      },
      "outputs": [],
      "source": [
        "def mse_loss(y_hat, y_target):\n",
        "    return (y_hat - y_target).pow(2.).mean()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XsqngFJO5HG0"
      },
      "source": [
        "The code cell below just initializes a couple of hyperparameters, and helper variables to log the progress of the trainig process."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wY1ZO4UJzkGV"
      },
      "outputs": [],
      "source": [
        "# learning rate\n",
        "lr = 0.1\n",
        "\n",
        "# maximum number of iterations\n",
        "iters = 100\n",
        "\n",
        "# used to log loss values for visualization purposes\n",
        "loss_values = []\n",
        "param_values = []"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pn6iZs6R5qKb"
      },
      "source": [
        "In this case, as the data is 2 dimensional the regressor will have two values: corresponding to the slope and intercept. These parameters will be stored in a single `Tensor` object initialized to random values:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1bjmquwy5QAQ"
      },
      "outputs": [],
      "source": [
        "# intialize the parameters of the model\n",
        "model = torch.randn(2, requires_grad=True)\n",
        "model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "49zGV9yE6T7e"
      },
      "source": [
        "This code block is where all the fun happens: *the training loop*.\n",
        "\n",
        "Step by step the training loop goes as the following:\n",
        "\n",
        "1. Given $x$, use the model to compute the estimate $\\hat{y}$.\n",
        "2. Compute the loss value based on the actual $y$ and the estimate $\\hat{y}$.\n",
        "3. Call `.backward()` on the loss value, this way autograd will fill the gradients of the tensors that produced $\\hat{y}$.\n",
        "4. Update the parameters of the model using gradient descent.\n",
        "5. Reset the gradients of the tensors to zero before the next iteration."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "adniGKBR2k3F"
      },
      "outputs": [],
      "source": [
        "for it in range(iters):\n",
        "\n",
        "    # 👁️ (1) compute the estimated y values using our model\n",
        "    y_hat = x * model[0] + model[1]\n",
        "\n",
        "    # 👁️ (2) compute the error\n",
        "    loss = mse_loss(y_hat, y)\n",
        "\n",
        "    # just for logging\n",
        "    loss_values.append(loss.item())\n",
        "    param_values.append(model.detach().tolist())\n",
        "    if it % 10 == 0:\n",
        "        print(f\"{it+1}/{    iters} MSE: {loss.item()}\")\n",
        "\n",
        "    # 👁️ (3) use autograd to compute gradients w.r.t the loss\n",
        "    loss.backward()\n",
        "\n",
        "    with torch.no_grad(): # within this block, autograd is disabled\n",
        "        # 👁️ (4) update the parameters of the model\n",
        "        # note that we don't want autograd to track this operation\n",
        "        model -= lr*model.grad\n",
        "\n",
        "    # 👁️ (5) set gradients to zero for the next iteration,\n",
        "    # otherwise gradients accumulate\n",
        "    model.grad.data.zero_()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oyb1P-Lb9G-M"
      },
      "source": [
        "Plot the data logged during the training process:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NAxMASqI89ZM"
      },
      "outputs": [],
      "source": [
        "# This is just an utility function, its content isn't relevant for the moment\n",
        "def viz_linear_reg(loss_values, param_values):\n",
        "    plt.subplot(2, 1, 1)\n",
        "    plt.plot(loss_values)\n",
        "    plt.title(\"Loss curve\")\n",
        "    plt.ylabel(\"MSE Loss\")\n",
        "    plt.xlabel(\"Iteration\")\n",
        "\n",
        "    plt.subplot(2, 1, 2)\n",
        "    plt.scatter(x, y, alpha=0.7)\n",
        "\n",
        "    for i, param in enumerate(param_values):\n",
        "        plt.axline((0, param[1]), slope=param[0],\n",
        "            color=[1., 0., 0., 0.5*i/len(param_values)])\n",
        "    plt.xlabel(\"X\")\n",
        "    plt.ylabel(\"Y\")\n",
        "    plt.title(\"Evolution of the paramaters\")\n",
        "    plt.tight_layout()\n",
        "\n",
        "viz_linear_reg(loss_values, param_values)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qllmuGLt-ijT"
      },
      "source": [
        "## A little bit more sophisticated example 🧞\n",
        "\n",
        "Despite the simplicity of the following example, we used quite \"low level\" constructs to conduct it. Normally, one doesn't directly define the parameters of the model as tensors, and doesn't perform gradient descent by hand.\n",
        "\n",
        "In this exercise, we'll apply the same concepts but in a higher level of abstraction, using the ergonomic functionalities that PyTorch offers.\n",
        "\n",
        "<br>\n",
        "\n",
        "First, let's initialize the same model but using PyTorch's [`nn`](https://pytorch.org/docs/stable/nn.html) module:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AXrnQH00_x-s"
      },
      "outputs": [],
      "source": [
        "import torch.nn as nn\n",
        "\n",
        "model = nn.Linear(1, 1)\n",
        "\n",
        "model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BTsvjrxA_e2x"
      },
      "source": [
        "Inspecting the `Linear` object (more info [here](https://pytorch.org/docs/stable/generated/torch.nn.Linear.html#linear)), we can see that it holds the same slope and intercept values as in the previous exercise."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5LFtlF6MAgjW"
      },
      "outputs": [],
      "source": [
        "model.weight"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jbYWJNfzA2YW"
      },
      "outputs": [],
      "source": [
        "model.bias"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kkSl0zfe_6d2"
      },
      "source": [
        "As promised, this time we won't perform gradient descent by hand, instead we'll employ the `SGD` optimizer from the PyTorch's [`optim`](https://pytorch.org/docs/stable/optim.html) module.\n",
        "\n",
        "Note that `SGD` is given the parameters of the model and the learning rate (optional) in the initialization. For the complete list of arguments and more info visit the [docs](https://pytorch.org/docs/stable/generated/torch.optim.SGD.html?highlight=sgd#torch.optim.SGD).  "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8pd0dPWj-1Cf"
      },
      "outputs": [],
      "source": [
        "from torch.optim import SGD\n",
        "\n",
        "optimizer = SGD(model.parameters(), lr=lr)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "01TsPXNNA8dF"
      },
      "source": [
        "The code cell below presents the modified training loop from the later exercise, in this case using PyTorch's `Linear` and `SGD` objects.\n",
        "Note that the loop continues to have the same five steps, however, in this case the optimization (step 4) is handled by PyTorch. Although basic SGD can be quite simple to implement by hand, it's handy to have more exotic (e.g. AdamW or RMSprop) optimizers efficiently implemented (although it's a fun exercise to do so)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wkSYQ9abBh6w"
      },
      "outputs": [],
      "source": [
        "loss_values = []\n",
        "param_values = []\n",
        "\n",
        "for it in range(100):\n",
        "    # 👁️ (1) compute the estimated y values using our model\n",
        "    y_hat = model(x)\n",
        "\n",
        "    # 👁️ (2) compute the error\n",
        "    loss = mse_loss(y_hat, y)\n",
        "\n",
        "\n",
        "    # ---- just for logging ---- #\n",
        "\n",
        "    loss_values.append(loss.item())\n",
        "    param_values.append([model.weight[0][0].item(),\n",
        "                         model.bias[0].item()])\n",
        "    if it % 10 == 0:\n",
        "        print(f\"{it+1}/{iters} MSE: {loss.item()}\")\n",
        "\n",
        "    # -------------------------- #\n",
        "\n",
        "\n",
        "    # 👁️ (3) tell autograd to compute gradients w.r.t the loss\n",
        "    loss.backward()\n",
        "\n",
        "    # 👁️ (4) update the parameters of the model\n",
        "    optimizer.step()\n",
        "\n",
        "    # 👁️ (5) clear previous gradient values\n",
        "    optimizer.zero_grad()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i5Pz7dyaDYpm"
      },
      "source": [
        "Visualize the training process:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WGKJesHBDK7n"
      },
      "outputs": [],
      "source": [
        "viz_linear_reg(loss_values, param_values)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vlk_0WAIncx-"
      },
      "source": [
        "## We are ready for the neurons 🧠  \n",
        "\n",
        "Up until this point we have seen the basics of tensors, their operations and properties, optimizers, and used autograd for linear regression. In this section, we'll use all this knowledge to train a neural network to classify MNIST digits!! 🔢 (I know... not super exciting, but a good stepping stone towards becoming a Pytorch ninja).\n",
        "\n",
        "Although this example might sound trivial, the described training process follows the same logic and concepts as many real-life use cases.  "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sM8sxK_KqOJ3"
      },
      "source": [
        "### Loading and preparing the data 📥\n",
        "\n",
        "The first step is to import `torchvision`, PyTorch's module for computer vision."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e-F58KZwqGJT"
      },
      "outputs": [],
      "source": [
        "import torchvision\n",
        "import torchvision.transforms as transforms"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-z7iJYdBLd-q"
      },
      "source": [
        "The following function is just an  utility function, and it's contents aren't very relevant for the objectives of the exercise. You can try to understand what the function does in depth later as an extra exercise."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EAWGswfQs7Q1"
      },
      "outputs": [],
      "source": [
        "# utility function\n",
        "def show_batch(batch):\n",
        "    batch = batch.view(-1, 1, 28, 28) # de-flatten\n",
        "    img = torchvision.utils.make_grid(batch)\n",
        "    img = img / 2 + 0.5 # unnormalize\n",
        "    npimg = img.numpy()\n",
        "    plt.imshow(np.transpose(npimg, (1, 2, 0)))\n",
        "    plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tDPdcgvGL7UL"
      },
      "source": [
        "Let's download and preprocess the MNIST dataset:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nJMKpGH6qcFV"
      },
      "outputs": [],
      "source": [
        "transform = transforms.Compose([\n",
        "     transforms.ToTensor(),\n",
        "     transforms.Normalize(0.5, 0.5),\n",
        "     lambda x: x.flatten(),\n",
        "    ])\n",
        "\n",
        "batch_size = 128\n",
        "\n",
        "# training data\n",
        "trainset = torchvision.datasets.MNIST(root='./data', train=True,\n",
        "                                        download=True, transform=transform)\n",
        "trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,\n",
        "                                          shuffle=True, num_workers=2)\n",
        "\n",
        "# test data\n",
        "testset = torchvision.datasets.MNIST(root='./data', train=False,\n",
        "                                       download=True, transform=transform)\n",
        "testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,\n",
        "                                         shuffle=False, num_workers=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J5Za-T5xMFTf"
      },
      "source": [
        "There's quite a lot going in the code cell above, let's go step by step:\n",
        "\n",
        "1. Define the transformation to apply to the dataset. Via `transforms.Compose` PyTorch enables to construct a pipeline of functions to preprocess the data before training (more info [here](https://pytorch.org/vision/main/transforms.html)). In this example, the pipeline goes as the following:\n",
        " 1. Images are first converted to tensors.\n",
        " 2. Normalize the images, as by default images are defined in the RGB space (values from 0 to 255).\n",
        " 3. MNIST contains 28x28 pixel images, flatten the images to 784 element vectors, as we'll consider a simple multi-layer perceptron and not a convolutional neural network.  \n",
        "\n",
        "2. Download the training set of MNIST and store it in a PyToch's `DataSet` object. This is used to download the data, load it to Python, and applying the `transform` defined in the step above. Then, prepare the dataset for training using the `DataLoader` object (more info on DataSet and DataLoaders [here](https://pytorch.org/tutorials/beginner/basics/data_tutorial.html)). The `DataLoader` is used to shuffle the dataset and split it in batches.\n",
        "\n",
        "3. Same for the test set of the MNIST dataset."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TBu5EmCwPkgs"
      },
      "source": [
        "Let's see how can we use the just created `DataLoader`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wQLnSfGLq9wY"
      },
      "outputs": [],
      "source": [
        "# convert the DataLoader to a python's iterator\n",
        "trainiter = iter(trainloader)\n",
        "\n",
        "# get the first element of the iterator (the first batch)\n",
        "images, labels = next(trainiter)\n",
        "\n",
        "print(f\"images: {images.size()}, labels: {labels.size()}\\n\\n\")\n",
        "\n",
        "show_batch(images[:4]) # show the first 4 images of the batch"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zyjOeWljtVN2"
      },
      "source": [
        "## Defining the model 🥼🧪\n",
        "\n",
        "Remember that in the previous exercises we have defined our model as a PyTorch `Tensor` and as a `Linear` module. As in this exercise we're going to build a more complex model (a MLP), it's convenint to define it as a PyTorch module (see the  [docs](https://pytorch.org/docs/stable/generated/torch.nn.Module.html)).\n",
        "PyTorch modules are objects normally used to compose models, where models can be composed by other modules (it's a bit of a matryoshka). In fact, `Linear` is an already defined module (there are many more, check [this](https://pytorch.org/docs/stable/nn.html)) that we can use to build our own one!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3ZPIt4_dteMf"
      },
      "outputs": [],
      "source": [
        "class Net(nn.Module): # inherit nn.Module\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "\n",
        "        # define the layers\n",
        "        self.l1 = nn.Linear(28*28, 512)\n",
        "        self.l2 = nn.Linear(512, 10)\n",
        "\n",
        "        # an activation function, feel free to change it\n",
        "        self.activation = nn.ReLU()\n",
        "\n",
        "        # as an extra, let's also add drop out (with 0.2 prob.)\n",
        "        self.d1 = nn.Dropout(0.2)\n",
        "\n",
        "    # this method will called every forward pass of our model\n",
        "    def forward(self, x):\n",
        "        # define the forward pass\n",
        "        x = self.l1(x)\n",
        "        x = self.activation(x)\n",
        "        x = self.d1(x)\n",
        "        x = self.l2(x)\n",
        "        return x"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i4JhFjpiS01G"
      },
      "source": [
        "Initialize our model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D-R_4ZfFtUu5"
      },
      "outputs": [],
      "source": [
        "model = Net().to(device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-fNOIghqS9k2"
      },
      "source": [
        "Note that `.to(device)` has been used to move the model to the selected device (\"cuda:0\" or \"cpu\" in our case)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a1-fkNMFvOlY"
      },
      "source": [
        "## The training loop 🔃\n",
        "\n",
        "As a common practice, we'll use a PyTorch optimizer to optimize our model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ahY8H61Nu0hh"
      },
      "outputs": [],
      "source": [
        "from torch.optim import Adam\n",
        "\n",
        "optimizer = Adam(model.parameters())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "axb_-NShTaIZ"
      },
      "source": [
        "In this case, for convenience, instead of implementing our own loss function we're going to use an already defined one, the [`CrossEntropyLoss`](https://pytorch.org/docs/2.0/generated/torch.nn.CrossEntropyLoss.html?highlight=crossentropyloss#torch.nn.CrossEntropyLoss). PyTorch implements many loss functions, they are available [here](https://pytorch.org/docs/stable/nn.html#loss-functions)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "naLO_0PkTaRf"
      },
      "outputs": [],
      "source": [
        "criterion = nn.CrossEntropyLoss()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2tid11S6URsx"
      },
      "source": [
        "And finally, here is the main training loop!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YUkUMXzVwFAk"
      },
      "outputs": [],
      "source": [
        "from tqdm import tqdm # used for the pretty progress bar\n",
        "\n",
        "num_epochs = 3\n",
        "\n",
        "for epoch in range(num_epochs): # for each epoch\n",
        "\n",
        "    loss_sum = 0 # used to compute the epoch's mean loss\n",
        "\n",
        "    for images, labels in tqdm(trainloader): # for each batch in the train set\n",
        "\n",
        "        # use the model to make a prediction\n",
        "        pred = model(images.to(device))\n",
        "\n",
        "        # compute the error of the model w.r.t the truth labels\n",
        "        loss = criterion(pred, labels.to(device))\n",
        "\n",
        "        loss_sum += loss.item() # used to compute the epoch's mean loss\n",
        "\n",
        "        # optimization\n",
        "        loss.backward()       # compute gradients\n",
        "        optimizer.step()      # run the optimizer\n",
        "        optimizer.zero_grad() # clear gradients for the next batch\n",
        "\n",
        "    print(f\"\\n** {epoch+1}/{num_epochs} mean loss: {loss_sum/len(trainloader)}\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "15F2e3g00X5C"
      },
      "source": [
        "## Evaluation 🔍\n",
        "\n",
        "Once we have our model trained, let's see how it behaves. The following code cell takes the first four images of the test set and uses the model to predict it's labels."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cPrY2OsE0xLf"
      },
      "outputs": [],
      "source": [
        "images, labels = next(iter(testloader)) # get the first batch\n",
        "\n",
        "# select the first four elements of the batch\n",
        "images = images[:4]\n",
        "labels = labels[:4]\n",
        "\n",
        "# plot the selected images\n",
        "show_batch(images)\n",
        "\n",
        "# When evaluating a model it must be set to eval mode, else dropout and other\n",
        "# regularization modules will still be active and bad things will happen.\n",
        "# To go back to training mode use model.train().\n",
        "model.eval() # ==> IMPORTANT!\n",
        "\n",
        "# use the model to make a prediction\n",
        "preds = model(images.to(device))    # note that the model outputs logits\n",
        "preds = torch.argmax(preds, dim=-1) # take the labels with greatest prob.\n",
        "print(f\"\\n** Predictions:\", preds.tolist())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qxI_HtKaYu1x"
      },
      "source": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DPES1_KfWl8m"
      },
      "source": [
        "<br>\n",
        "<br>\n",
        "\n",
        "And this is all! Congrats for reaching here!!👏🪘🪇 You're one step closer to\n",
        "become a PyTorch ninja🥷, but that's a journey you will have to follow alone 🎒🏕🏋🏞🗺.\n",
        "\n",
        "Hope you liked the tutorial and learned something!\n",
        "\n",
        "Agur!! 👋🤓"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "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.11.2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}