pytorch_basics/pytorch_NeuralNetworks.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/makagan/SSI_Projects/blob/main/pytorch_basics/pytorch_NeuralNetworks.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LYRTcMsjffTd"
   },
   "source": [
    "### Basics of ML with Pytorch\n",
    "\n",
    "Adapted from official [Pytorch tutorial](https://pytorch.org/tutorials/beginner/basics/intro.html) for dealing with PyTorch tensors, datasets, building neural networks etc., also has an accompanying video series."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zLqa3_Xp71et"
   },
   "source": [
    "## Build the Neural Network\n",
    "\n",
    "Neural networks comprise of layers/modules that perform operations on data. The [torch.nn](https://pytorch.org/docs/stable/nn.html) namespace provides all the building blocks you need to build your own neural network. Every module in PyTorch subclasses the [nn.Module](https://pytorch.org/docs/stable/generated/torch.nn.Module.html). A neural network is a module itself that consists of other modules (layers). This nested structure allows for building and managing complex architectures easily.\n",
    "\n",
    "In the following sections, we’ll build a neural network to classify images in the FashionMNIST dataset.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WmkoCs_d8QXU"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets, transforms\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PYVHRDnm8TOi"
   },
   "source": [
    "### Get Device for Training\n",
    "\n",
    "We want to be able to train our model on a hardware accelerator like the GPU, if it is available. Let’s check to see if [torch.cuda](https://pytorch.org/docs/stable/notes/cuda.html) is available, else we continue to use the CPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lSsakf_z8bZo"
   },
   "outputs": [],
   "source": [
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(f\"Using {device} device\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qLhOQc6Y8fPh"
   },
   "source": [
    "### Define the Class\n",
    "\n",
    "We define our neural network by subclassing `nn.Module`, and initialize the neural network layers in `__init__`. Every `nn.Module` subclass implements the operations on input data in the `forward` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "A1JzqTbt8kjO"
   },
   "outputs": [],
   "source": [
    "class NeuralNetwork(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(NeuralNetwork, self).__init__()\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.linear_relu_stack = nn.Sequential(\n",
    "            nn.Linear(28*28, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 10),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.flatten(x)\n",
    "        logits = self.linear_relu_stack(x)\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "J1T4HVeE8po9"
   },
   "source": [
    "We create an instance of `NeuralNetwork`, and move it to the `device`, and print its structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "i5S2Y0UX8t5D"
   },
   "outputs": [],
   "source": [
    "model = NeuralNetwork().to(device)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iDYias168yR1"
   },
   "source": [
    "To use the model, we pass it the input data. This executes the model’s `forward`.\n",
    "\n",
    "Calling the model on the input returns a 2-dimensional tensor with dim=0 corresponding to each output of 10 raw predicted values for each class, and dim=1 corresponding to the individual values of each output. We get the prediction probabilities by passing it through an instance of the `nn.Softmax` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Yjj-R1zi9f9_"
   },
   "outputs": [],
   "source": [
    "X = torch.rand(1, 28, 28, device=device)\n",
    "logits = model(X)\n",
    "pred_probab = nn.Softmax(dim=1)(logits)\n",
    "y_pred = pred_probab.argmax(1)\n",
    "print(f\"Predicted class: {y_pred}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wGDMGpIR-Uvl"
   },
   "source": [
    "### Model Layers\n",
    "\n",
    "Let’s break down the layers in the FashionMNIST model. To illustrate it, we will take a sample minibatch of 3 images of size 28x28 and see what happens to it as we pass it through the network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4iz2NGn7-cgB"
   },
   "outputs": [],
   "source": [
    "input_image = torch.rand(3,28,28)\n",
    "print(input_image.size())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "g5MvoDWG-gtz"
   },
   "source": [
    "**nn.Flatten**\n",
    "\n",
    "We initialize the [`nn.Flatten`](https://pytorch.org/docs/stable/generated/torch.nn.Flatten.html) layer to convert each 2D 28x28 image into a contiguous array of 784 pixel values (the minibatch dimension (at dim=0) is maintained).\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nPCEX-2V-4pI"
   },
   "outputs": [],
   "source": [
    "flatten = nn.Flatten()\n",
    "flat_image = flatten(input_image)\n",
    "print(flat_image.size())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hYX_wGvz-31r"
   },
   "source": [
    "**nn.Linear**\n",
    "\n",
    "The [linear layer](https://pytorch.org/docs/stable/generated/torch.nn.Linear.html) is a module that applies a linear transformation on the input using its stored weights and biases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0OTtNFni-7uZ"
   },
   "outputs": [],
   "source": [
    "layer1 = nn.Linear(in_features=28*28, out_features=20)\n",
    "hidden1 = layer1(flat_image)\n",
    "print(hidden1.size())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8CkWAD1B_WN1"
   },
   "source": [
    "**nn.ReLU**\n",
    "\n",
    "Non-linear activations are what create the complex mappings between the model’s inputs and outputs. They are applied after linear transformations to introduce nonlinearity, helping neural networks learn a wide variety of phenomena.\n",
    "\n",
    "In this model, we use [`nn.ReLU`](https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html) between our linear layers, but there’s other activations to introduce non-linearity in your model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3kW1AbfK_kc3"
   },
   "outputs": [],
   "source": [
    "print(f\"Before ReLU: {hidden1}\\n\\n\")\n",
    "hidden1 = nn.ReLU()(hidden1)\n",
    "print(f\"After ReLU: {hidden1}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "x8trQ1l6_sz-"
   },
   "source": [
    "**nn.Sequential**\n",
    "\n",
    "[`nn.Sequential`](https://pytorch.org/docs/stable/generated/torch.nn.Sequential.html) is an ordered container of modules. The data is passed through all the modules in the same order as defined. You can use sequential containers to put together a quick network like `seq_modules`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "G6UpMoKCAM1t"
   },
   "outputs": [],
   "source": [
    "seq_modules = nn.Sequential(\n",
    "    flatten,\n",
    "    layer1,\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(20, 10)\n",
    ")\n",
    "input_image = torch.rand(3,28,28)\n",
    "logits = seq_modules(input_image)\n",
    "print(logits)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5FJ-3DB-ASqc"
   },
   "source": [
    "**nn.Softmax**\n",
    "\n",
    "The last linear layer of the neural network returns logits - values in [-infty, infty] - which are passed to the [`nn.Softmax`](https://pytorch.org/docs/stable/generated/torch.nn.Softmax.html) module. The logits are scaled to values [0, 1] representing the model’s predicted probabilities for each class. `dim` parameter indicates the dimension along which the values must sum to 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_wxcRRgiAjza"
   },
   "outputs": [],
   "source": [
    "softmax = nn.Softmax(dim=1) #softmax is over the columns per each row\n",
    "pred_probab = softmax(logits)\n",
    "print(pred_probab)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "glE-gs_3H2i3"
   },
   "source": [
    "### Model Parameters\n",
    "\n",
    "Many layers inside a neural network are parameterized, i.e. have associated weights and biases that are optimized during training. Subclassing `nn.Module` automatically tracks all fields defined inside your model object, and makes all parameters accessible using your model’s `parameters()` or `named_parameters()` methods.\n",
    "\n",
    "In this example, we iterate over each parameter, and print its size and a preview of its values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pvAKVy-jJhzY"
   },
   "outputs": [],
   "source": [
    "print(f\"Model structure: {model}\\n\\n\")\n",
    "\n",
    "for name, param in model.named_parameters():\n",
    "    print(f\"Layer: {name} | Size: {param.size()} | Values : {param[:2]} \\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qnGZJmo0JyFJ"
   },
   "source": [
    "## Automatic Differentiation with `torch.autograd`\n",
    "\n",
    "When training neural networks, the most frequently used algorithm is **back propagation**. In this algorithm, parameters (model weights) are adjusted according to the **gradient** of the loss function with respect to the given parameter.\n",
    "\n",
    "To compute those gradients, PyTorch has a built-in differentiation engine called `torch.autograd`. It supports automatic computation of gradient for any computational graph.\n",
    "\n",
    "Consider the simplest one-layer neural network, with input `x`, parameters `w` and `b`, and some loss function. It can be defined in PyTorch in the following manner:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "86v8CQO4-eM7"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "x = torch.ones(5)  # input tensor\n",
    "y = torch.zeros(3)  # expected output\n",
    "w = torch.randn(5, 3, requires_grad=True)\n",
    "b = torch.randn(3, requires_grad=True)\n",
    "z = torch.matmul(x, w)+b\n",
    "loss = torch.nn.functional.binary_cross_entropy_with_logits(z, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "seY2qP4pgpf5"
   },
   "source": [
    "In this network, `w` and `b` are parameters, which we need to optimize. Thus, we need to be able to compute the gradients of loss function with respect to those variables. In order to do that, we set the `requires_grad` property of those tensors."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hQzbakYZh7nz"
   },
   "source": [
    "### Computing Gradients\n",
    "\n",
    "To optimize weights of parameters in the neural network, we need to compute the derivatives of our loss function with respect to the parameters under some fixed values of `x` and `y`. To compute those derivatives, we call `loss.backward()`, and then retrieve the values from `w.grad` and `b.grad`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1l_xiMb5hCoR"
   },
   "outputs": [],
   "source": [
    "loss.backward()\n",
    "print(w.grad)\n",
    "print(b.grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TI9PXO7BFP1W"
   },
   "source": [
    "**NOTE**\n",
    "\n",
    "\n",
    "\n",
    "* We can only obtain the grad properties for the leaf nodes of the computational graph, which have `requires_grad` property set to `True`. For all other nodes in our graph, gradients will not be available.\n",
    "\n",
    "* We can only perform gradient calculations using `backward` once on a given graph, for performance reasons. If we need to do several backward calls on the same graph, we need to pass `retain_graph=True` to the backward call.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AvlUPDwCjgDu"
   },
   "source": [
    "### Disabling Gradient Tracking\n",
    "\n",
    "By default, all tensors with `requires_grad=True` are tracking their computational history and support gradient computation. However, there are some cases when we do not need to do that, for example, when we have trained the model and just want to apply it to some input data, i.e. we only want to do forward computations through the network. We can stop tracking computations by surrounding our computation code with `torch.no_grad()` block:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Rx4qDH0ljmR5"
   },
   "outputs": [],
   "source": [
    "z = torch.matmul(x, w)+b\n",
    "print(z.requires_grad)\n",
    "\n",
    "with torch.no_grad():\n",
    "    z = torch.matmul(x, w)+b\n",
    "print(z.requires_grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KvlFBn4cjqqF"
   },
   "source": [
    "Another way to achieve the same result is to use the `detach()` method on the tensor:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "n_BQ-pjFjsBF"
   },
   "outputs": [],
   "source": [
    "z = torch.matmul(x, w)+b\n",
    "z_det = z.detach()\n",
    "print(z_det.requires_grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QFm9kAFxjw6b"
   },
   "source": [
    "\n",
    "There are reasons you might want to disable gradient tracking:\n",
    "\n",
    "* To mark some parameters in your neural network as **frozen parameters**.\n",
    "\n",
    "* To **speed up computations** when you are only doing forward pass, because computations on tensors that do not track gradients would be more efficient.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Sr85ZJ-Clthh"
   },
   "source": [
    "### Computational graphs\n",
    "\n",
    "The simple NN created above defines the following **computational graph**:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hf9ON63LmYMm"
   },
   "source": [
    "![](https://pytorch.org/tutorials/_images/comp-graph.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aobf11G1mZQq"
   },
   "source": [
    "A function that we apply to tensors to construct computational graph is in fact an object of class `Function`. This object knows how to compute the function in the *forward* direction, and also how to compute its derivative during the *backward* propagation step. A reference to the backward propagation function is stored in `grad_fn` property of a tensor. You can find more information of `Function` [in the documentation](https://pytorch.org/docs/stable/autograd.html#function)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2GDihJRdnBWh"
   },
   "outputs": [],
   "source": [
    "x = torch.ones(5)  # input tensor\n",
    "y = torch.zeros(3)  # expected output\n",
    "w = torch.randn(5, 3, requires_grad=True)\n",
    "b = torch.randn(3, requires_grad=True)\n",
    "z = torch.matmul(x, w)+b\n",
    "loss = torch.nn.functional.binary_cross_entropy_with_logits(z, y)\n",
    "\n",
    "print(f\"Gradient function for z = {z.grad_fn}\")\n",
    "print(f\"Gradient function for loss = {loss.grad_fn}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MbbBxbm_nsh1"
   },
   "source": [
    "Conceptually, autograd keeps a record of data (tensors) and all executed operations (along with the resulting new tensors) in a [directed acyclic graph (DAG)](https://en.wikipedia.org/wiki/Directed_acyclic_graph) consisting of [`Function`](https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function) objects. In this DAG, leaves are the input tensors, roots are the output tensors. By tracing this graph from roots to leaves, you can automatically compute the gradients using the chain rule.\n",
    "\n",
    "In a forward pass, autograd does two things simultaneously:\n",
    "\n",
    "* run the requested operation to compute a resulting tensor\n",
    "* maintain the operation’s gradient function in the DAG.\n",
    "\n",
    "The backward pass kicks off when `.backward()` is called on the DAG root. autograd then:\n",
    "\n",
    "* computes the gradients from each `.grad_fn`,\n",
    "* accumulates them in the respective tensor’s `.grad` attribute\n",
    "* using the chain rule, propagates all the way to the leaf tensors.\n",
    "\n",
    "**NOTE**\n",
    "\n",
    "DAGs are dynamic in PyTorch An important thing to note is that the graph is recreated from scratch; after each `.backward()` call, autograd starts populating a new graph. This is exactly what allows you to use control flow statements in your model; you can change the shape, size and operations at every iteration if needed.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bfu7n8KiwlLO"
   },
   "source": [
    "### Tensor Gradients and Jacobian Products\n",
    "\n",
    "In many cases, we have a scalar loss function, and we need to compute the gradient with respect to some parameters. However, there are cases when the output function is an arbitrary tensor. In this case, PyTorch allows you to compute so-called **Jacobian product**, and not the actual gradient.\n",
    "\n",
    "For a vector function $\\vec{y}=f(\\vec{x})$, where $\\vec{x} = \\langle x_1,...,x_n \\rangle$ and $\\vec{y} = \\langle y_1,...,y_m \\rangle$, a gradient of $\\vec{y}$ with respect to $\\vec{x}$ is given by the **Jacobian matrix**: \n",
    "\n",
    "\\begin{equation}\n",
    "J = \\begin{pmatrix}\n",
    "\\frac{\\partial y_1}{\\partial x_1} & ... & \\frac{\\partial y_1}{\\partial x_n}\\\\\n",
    "\\vdots & \\ddots & \\vdots\\\\\n",
    "\\frac{\\partial y_m}{\\partial x_1} & ... & \\frac{\\partial y_n}{\\partial x_n}\n",
    "\\end{pmatrix}\n",
    "\\end{equation}\n",
    "\n",
    "Instead of computing the Jacobian matrix itself, PyTorch allows you to compute **Jacobian Product** $v^T \\cdot J$ for the given input vector $\\vec{v} = \\langle v_1,...,v_m \\rangle$. This is achieved by calling `backward` with $v$ as an argument. The size of $\\vec{v}$ should be the same as the size of the original tensor, with respect to which we want to compute the product:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "P5XDfCqtnsKW"
   },
   "outputs": [],
   "source": [
    "inp = torch.eye(4, 5, requires_grad=True) #Returns a 2-D tensor with ones on the diagonal and zeros elsewhere.\n",
    "out = (inp+1).pow(2).t()\n",
    "out.backward(torch.ones_like(out), retain_graph=True)\n",
    "print(f\"First call\\n{inp.grad}\")\n",
    "out.backward(torch.ones_like(out), retain_graph=True)\n",
    "print(f\"\\nSecond call\\n{inp.grad}\")\n",
    "inp.grad.zero_()\n",
    "out.backward(torch.ones_like(out), retain_graph=True)\n",
    "print(f\"\\nCall after zeroing gradients\\n{inp.grad}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XlZ_U19GJuK3"
   },
   "source": [
    "Notice that when we call `backward` for the second time with the same argument, the value of the gradient is different. This happens because when doing backward propagation, PyTorch accumulates the gradients, i.e. the value of computed gradients is added to the grad property of all leaf nodes of computational graph. If you want to compute the proper gradients, you need to zero out the grad property before. In real-life training an optimizer helps us to do this.\n",
    "\n",
    "**NOTE**\n",
    "\n",
    "Previously we were calling the `backward()` function without parameters. This is essentially equivalent to calling `backward(torch.tensor(1.0))`, which is a useful way to compute the gradients in case of a scalar-valued function, such as loss during neural network training."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "E5tlCwRvKh0g"
   },
   "source": [
    "## Optimizing Model Parameters\n",
    "\n",
    "Now that we have a model and data it’s time to train, validate and test our model by optimizing its parameters on our data. Training a model is an iterative process; in each iteration the model makes a guess about the output, calculates the error in its guess (loss), collects the derivatives of the error with respect to its parameters, and optimizes these parameters using gradient descent. \n",
    "\n",
    "Let's start from scratch:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Tt7COBEhK79J"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets\n",
    "from torchvision.transforms import ToTensor\n",
    "\n",
    "training_data = datasets.FashionMNIST(\n",
    "    root=\"data\",\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=ToTensor()\n",
    ")\n",
    "\n",
    "test_data = datasets.FashionMNIST(\n",
    "    root=\"data\",\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=ToTensor()\n",
    ")\n",
    "\n",
    "train_dataloader = DataLoader(training_data, batch_size=64)\n",
    "test_dataloader = DataLoader(test_data, batch_size=64)\n",
    "\n",
    "class NeuralNetwork(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(NeuralNetwork, self).__init__()\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.linear_relu_stack = nn.Sequential(\n",
    "            nn.Linear(28*28, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 10),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.flatten(x)\n",
    "        logits = self.linear_relu_stack(x)\n",
    "        return logits\n",
    "\n",
    "model = NeuralNetwork()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "03gM7g6SLVa3"
   },
   "source": [
    "### Hyperparameters\n",
    "\n",
    "Hyperparameters are adjustable parameters that let you control the model optimization process. Different hyperparameter values can impact model training and convergence rates (read more about hyperparameter tuning)\n",
    "\n",
    "We define the following hyperparameters for training:\n",
    "\n",
    "* **Number of Epochs** - the number times to iterate over the dataset\n",
    "* **Batch Size** - the number of data samples propagated through the network before the parameters are updated\n",
    "* **Learning Rate** - how much to update models parameters at each batch/epoch. Smaller values yield slow learning speed, while large values may result in unpredictable behavior during training.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YhsiRwTGLdrY"
   },
   "outputs": [],
   "source": [
    "learning_rate = 1e-3\n",
    "batch_size = 64\n",
    "epochs = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eLaYx8OpLg_A"
   },
   "source": [
    "### Optimization Loop\n",
    "\n",
    "Once we set our hyperparameters, we can then train and optimize our model with an optimization loop. Each iteration of the optimization loop is called an epoch.\n",
    "\n",
    "Each epoch consists of two main parts:\n",
    "\n",
    "* **Train Loop**- iterate over the training dataset and try to converge to optimal parameters.\n",
    "* **Validation/Test Loop** - iterate over the test dataset to check if model performance is improving.\n",
    "\n",
    "Let’s briefly familiarize ourselves with some of the concepts used in the training loop."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9n_BnszFLwFo"
   },
   "source": [
    "### Loss function\n",
    "\n",
    "When presented with some training data, our untrained network is likely not to give the correct answer. Loss function measures the degree of dissimilarity of obtained result to the target value, and it is the loss function that we want to minimize during training. To calculate the loss we make a prediction using the inputs of our given data sample and compare it against the true data label value.\n",
    "\n",
    "Common loss functions include [`nn.MSELoss`](https://pytorch.org/docs/stable/generated/torch.nn.MSELoss.html#torch.nn.MSELoss) (Mean Square Error) for regression tasks, and [`nn.NLLLoss`](https://pytorch.org/docs/stable/generated/torch.nn.NLLLoss.html#torch.nn.NLLLoss) (Negative Log Likelihood) for classification. [`nn.CrossEntropyLoss`](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html#torch.nn.CrossEntropyLoss) combines `nn.LogSoftmax` and `nn.NLLLoss`.\n",
    "\n",
    "We pass our model’s output logits to `nn.CrossEntropyLoss`, which will normalize the logits and compute the prediction error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6EfUbtr7LswF"
   },
   "outputs": [],
   "source": [
    "# Initialize the loss function\n",
    "loss_fn = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zjKDXHR8MXXW"
   },
   "source": [
    "### Optimizer\n",
    "\n",
    "Optimization is the process of adjusting model parameters to reduce model error in each training step. Optimization algorithms define how this process is performed (in this example we use Stochastic Gradient Descent). All optimization logic is encapsulated in the `optimizer` object. Here, we use the SGD optimizer; additionally, there are many [different optimizers](https://pytorch.org/docs/stable/optim.html) available in PyTorch such as ADAM and RMSProp, that work better for different kinds of models and data.\n",
    "\n",
    "We initialize the optimizer by registering the model’s parameters that need to be trained, and passing in the learning rate hyperparameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9fSFqrYXMjvs"
   },
   "outputs": [],
   "source": [
    "optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Po1-t3c5MoRr"
   },
   "source": [
    "Inside the training loop, optimization happens in three steps:\n",
    "\n",
    "* Call `optimizer.zero_grad()` to reset the gradients of model parameters. Gradients by default add up; to prevent double-counting, we explicitly zero them at each iteration.\n",
    "* Backpropagate the prediction loss with a call to `loss.backward()`. PyTorch deposits the gradients of the loss w.r.t. each parameter.\n",
    "* Once we have our gradients, we call `optimizer.step()` to adjust the parameters by the gradients collected in the backward pass.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "r0xYC5L7Mym-"
   },
   "source": [
    "### Full Implementation\n",
    "\n",
    "We define `train_loop` that loops over our optimization code, and `test_loop` that evaluates the model’s performance against our test data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pqTiQEXHgaQW"
   },
   "outputs": [],
   "source": [
    "def train_loop(dataloader, model, loss_fn, optimizer):\n",
    "    size = len(dataloader.dataset)\n",
    "    for batch, (X, y) in enumerate(dataloader):\n",
    "        # Compute prediction and loss\n",
    "        pred = model(X)\n",
    "        loss = loss_fn(pred, y)\n",
    "\n",
    "        # Backpropagation\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        if batch % 100 == 0:\n",
    "            loss, current = loss.item(), (batch + 1) * len(X)\n",
    "            print(f\"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]\")\n",
    "\n",
    "\n",
    "def test_loop(dataloader, model, loss_fn):\n",
    "    size = len(dataloader.dataset)\n",
    "    num_batches = len(dataloader)\n",
    "    test_loss, correct = 0, 0\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for X, y in dataloader:\n",
    "            pred = model(X)\n",
    "            test_loss += loss_fn(pred, y).item()\n",
    "            correct += (pred.argmax(1) == y).type(torch.float).sum().item()\n",
    "\n",
    "    test_loss /= num_batches\n",
    "    correct /= size\n",
    "    print(f\"Test Error: \\n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "q93MwM6OM7Zh"
   },
   "source": [
    "We initialize the loss function and optimizer, and pass it to `train_loop` and `test_loop`. Feel free to increase the number of epochs to track the model’s improving performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "moHSkwEJM_n-"
   },
   "outputs": [],
   "source": [
    "loss_fn = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)\n",
    "\n",
    "for t in range(epochs):\n",
    "    print(f\"Epoch {t+1}\\n-------------------------------\")\n",
    "    train_loop(train_dataloader, model, loss_fn, optimizer)\n",
    "    test_loop(test_dataloader, model, loss_fn)\n",
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gBkiZ4TlNENP"
   },
   "source": [
    "## Save and Load the Model\n",
    "\n",
    "In this section we will look at how to persist model state with saving, loading and running model predictions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rt9B8MHcNO5o"
   },
   "source": [
    "### Saving and Loading Model Weights\n",
    "\n",
    "PyTorch models store the learned parameters in an internal state dictionary, called `state_dict`. These can be persisted via the `torch.save` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HBdz4S6gNVeh"
   },
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), 'model_weights.pth')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7ZiO4p2sOa0i"
   },
   "source": [
    "To load model weights, you need to create an instance of the same model first, and then load the parameters using `load_state_dict()` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QkM47JCzOJff"
   },
   "outputs": [],
   "source": [
    "model_new = NeuralNetwork()\n",
    "model_new.load_state_dict(torch.load('model_weights.pth'))\n",
    "model_new.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yg5vIJHAOmpZ"
   },
   "source": [
    "### Saving and Loading Models with Shapes\n",
    "\n",
    "When loading model weights, we needed to instantiate the model class first, because the class defines the structure of a network. We might want to save the structure of this class together with the model, in which case we can pass model (and not `model.state_dict()`) to the saving function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CInluGULO90H"
   },
   "outputs": [],
   "source": [
    "torch.save(model, 'model.pth')\n",
    "model_new = torch.load('model.pth')"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "authorship_tag": "ABX9TyOcVpPAUZc0XPgH5j9zG3I3",
   "include_colab_link": true,
   "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.8.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}