CS330_PyTorch_Tutorial.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/DrSquare/AI_Coding/blob/main/CS330_PyTorch_Tutorial.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FlkUWhvSu-ja"
      },
      "source": [
        "Author: Kyle Hsu\n",
        "\n",
        "With revisions from: Rafael Rafailov, Evan Liu, Fahim Tajwar"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HI43aS5mRO1s"
      },
      "source": [
        "# Introduction\n",
        "Welcome to the PyTorch tutorial for CS330! This colab notebook accompanies [these slides](https://docs.google.com/presentation/d/1e_md1C24vZsMNtbAJyNyDkau99owklm65IG9MudizvE/edit?usp=sharing). If you haven't already, enable a GPU for this colab instance by doing \"Edit\" -> \"Notebook settings\" -> \"Hardware accelerator\" drop-down -> \"GPU\" -> \"Save\"."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H7cLA7izR_zl"
      },
      "source": [
        "Let's make sure we're using the right Python and PyTorch versions, and that we have a GPU at our disposal."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "I12wdApTZC_H",
        "outputId": "0e9b560e-6354-497d-e2fe-18f596786e29"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Python version info: 3.11.11 (main, Dec  4 2024, 08:55:07) [GCC 11.4.0]\n",
            "PyTorch version info: 2.5.1+cu124\n",
            "PyTorch detects a GPU: True\n"
          ]
        }
      ],
      "source": [
        "import sys\n",
        "import torch\n",
        "print(f'Python version info: {sys.version}')\n",
        "print(f'PyTorch version info: {torch.__version__}')\n",
        "print(f'PyTorch detects a GPU: {torch.cuda.is_available()}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ptKO3cRFTjWo"
      },
      "source": [
        "# Overview\n",
        "We'll begin by doing some MNIST classification, starting from low-level operations and gradually replacing chunks using `torch` abstractions such as `torch.nn.Module`, `torch.optim.SGD`, and `torch.utils.data.DataLoader`. (Much of this content is adapted from [this tutorial](https://pytorch.org/tutorials/beginner/nn_tutorial.html) by Jeremy Howard.) We'll then proceed to a whirlwind tour over PyTorch features that we suspect will be relevant for your homeworks for this course. Finally, we'll see some tips and tricks for debugging and getting help."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8t8qKU2yVE5Q"
      },
      "source": [
        "# PyTorch Basics via MNIST\n",
        "### MNIST\n",
        "Download MNIST."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "4VB4quTfVHCb"
      },
      "outputs": [],
      "source": [
        "from pathlib import Path\n",
        "import requests\n",
        "\n",
        "DATA_PATH = Path(\"./data\")\n",
        "PATH = DATA_PATH / \"mnist\"\n",
        "\n",
        "PATH.mkdir(parents=True, exist_ok=True)\n",
        "\n",
        "URL = \"https://github.com/pytorch/tutorials/raw/main/_static/\"\n",
        "FILENAME = \"mnist.pkl.gz\"\n",
        "\n",
        "if not (PATH / FILENAME).exists():\n",
        "        content = requests.get(URL + FILENAME).content\n",
        "        (PATH / FILENAME).open(\"wb\").write(content)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G7K4tLsIVURi"
      },
      "source": [
        "Load MNIST into memory."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "EzxqdhhKVTMA"
      },
      "outputs": [],
      "source": [
        "import pickle\n",
        "import gzip\n",
        "import numpy as np\n",
        "\n",
        "with gzip.open((PATH / FILENAME).as_posix(), \"rb\") as f:\n",
        "        ((x_train, y_train), (x_val, y_val), _) = pickle.load(f, encoding=\"latin-1\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ouKpbI6DWYWx"
      },
      "source": [
        "The data currently exists as NumPy arrays. Let's take a closer look."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 517
        },
        "id": "EG9KpubhWbNZ",
        "outputId": "cd7c50c6-1e06-448e-c823-66c00ebdad8f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "training inputs: shape (50000, 784), dtype float32\n",
            "training outputs: shape (50000,), dtype int64\n",
            "input range: (0.0, 0.99609375)\n",
            "label range: (0, 9)\n",
            "label of training example 42: 7\n"
          ]
        },
        {
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAaAAAAGdCAYAAABU0qcqAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlHJYcgAAAAlwSFlzAAAPYQAAD2EBqD+naQAAGhhJREFUeJzt3X9MVff9x/HX9QdX28J1iHC5Ey3aVpv6o5lTSmydjQRhifHnoq1/6GI0OmymtGvD0mp1S9hc0pnu6+yyZLo21XZuVatZXBQLxg1t/BVnthIhdOgU/LFwL6Kikc/3D9O73frz4L284fp8JCeRe8+H++7prc8eOBx8zjknAAA6WQ/rAQAADycCBAAwQYAAACYIEADABAECAJggQAAAEwQIAGCCAAEATPSyHuDr2tvbdebMGaWmpsrn81mPAwDwyDmnlpYWhUIh9ehx5/OcLhegM2fOKCcnx3oMAMADOnXqlAYOHHjH57vcl+BSU1OtRwAAxMG9/j5PWIDWrVunxx9/XH369FFeXp4+//zz+1rHl90AIDnc6+/zhATo448/VmlpqVauXKkjR45o9OjRmjx5ss6dO5eIlwMAdEcuAcaNG+dKSkqiH9+4ccOFQiFXXl5+z7XhcNhJYmNjY2Pr5ls4HL7r3/dxPwO6du2aDh8+rIKCguhjPXr0UEFBgaqrq2/Zv62tTZFIJGYDACS/uAfowoULunHjhrKysmIez8rKUmNj4y37l5eXKxAIRDeugAOAh4P5VXBlZWUKh8PR7dSpU9YjAQA6Qdx/DigjI0M9e/ZUU1NTzONNTU0KBoO37O/3++X3++M9BgCgi4v7GVBKSorGjBmjioqK6GPt7e2qqKhQfn5+vF8OANBNJeROCKWlpZo3b56+/e1va9y4cVq7dq1aW1v1/e9/PxEvBwDohhISoNmzZ+v8+fNasWKFGhsb9eyzz2rXrl23XJgAAHh4+ZxzznqI/xWJRBQIBKzHAAA8oHA4rLS0tDs+b34VHADg4USAAAAmCBAAwAQBAgCYIEAAABMECABgggABAEwQIACACQIEADBBgAAAJggQAMAEAQIAmCBAAAATBAgAYIIAAQBMECAAgAkCBAAwQYAAACYIEADABAECAJggQAAAEwQIAGCCAAEATBAgAIAJAgQAMEGAAAAmCBAAwAQBAgCYIEAAABMECABgggABAEwQIACACQIEADBBgAAAJggQAMAEAQIAmCBAAAATBAgAYIIAAQBMECAAgAkCBAAwQYAAACYIEADABAECAJggQAAAEwQIAGCCAAEATBAgAIAJAgQAMEGAAAAmCBAAwAQBAgCYIEAAABMECABgggABAEzEPUBvv/22fD5fzDZ8+PB4vwwAoJvrlYhP+swzz2jPnj3/fZFeCXkZAEA3lpAy9OrVS8FgMBGfGgCQJBLyPaCTJ08qFAppyJAhmjt3rhoaGu64b1tbmyKRSMwGAEh+cQ9QXl6eNm7cqF27dmn9+vWqr6/XCy+8oJaWltvuX15erkAgEN1ycnLiPRIAoAvyOedcIl+gublZgwcP1jvvvKMFCxbc8nxbW5va2tqiH0ciESIEAEkgHA4rLS3tjs8n/OqAfv366amnnlJtbe1tn/f7/fL7/YkeAwDQxST854AuXbqkuro6ZWdnJ/qlAADdSNwD9Nprr6mqqkpffvml/va3v2n69Onq2bOnXnrppXi/FACgG4v7l+BOnz6tl156SRcvXtSAAQP0/PPP68CBAxowYEC8XwoA0I0l/CIEryKRiAKBgPUYAIAHdK+LELgXHADABAECAJggQAAAEwQIAGCCAAEATBAgAIAJAgQAMEGAAAAmCBAAwAQBAgCYIEAAABMECABgIuG/kA5Idj6fz/Oajvx+rO9973ue18yaNcvzGkkaOnSo5zXPPfec5zUNDQ2e1yB5cAYEADBBgAAAJggQAMAEAQIAmCBAAAATBAgAYIIAAQBMECAAgAkCBAAwQYAAACYIEADABAECAJggQAAAE9wNG0lp4MCBHVo3depUz2vmzJnjec348eM9r+lMra2tntdcvnw5AZMgmXEGBAAwQYAAACYIEADABAECAJggQAAAEwQIAGCCAAEATBAgAIAJAgQAMEGAAAAmCBAAwAQBAgCY4Gak6FSjRo3yvKasrMzzmunTp3teI0kpKSme13z55Zee1/zf//2f5zW9enn/z3Xx4sWe10jS7t27Pa+5cOFCh14LDy/OgAAAJggQAMAEAQIAmCBAAAATBAgAYIIAAQBMECAAgAkCBAAwQYAAACYIEADABAECAJggQAAAE9yMFHrxxRc7tO53v/ud5zVZWVme1/Tp08fzmt/+9ree10jSBx984HnNkSNHPK+5fPmy5zXPPvus5zUdvRnp3//+9w6tA7zgDAgAYIIAAQBMeA7Qvn37NGXKFIVCIfl8Pm3bti3meeecVqxYoezsbPXt21cFBQU6efJkvOYFACQJzwFqbW3V6NGjtW7duts+v2bNGr377rt67733dPDgQT366KOaPHmyrl69+sDDAgCSh+eLEIqLi1VcXHzb55xzWrt2rd58801NnTpVkvT+++8rKytL27Zt05w5cx5sWgBA0ojr94Dq6+vV2NiogoKC6GOBQEB5eXmqrq6+7Zq2tjZFIpGYDQCQ/OIaoMbGRkm3XmqblZUVfe7rysvLFQgEoltOTk48RwIAdFHmV8GVlZUpHA5Ht1OnTlmPBADoBHENUDAYlCQ1NTXFPN7U1BR97uv8fr/S0tJiNgBA8otrgHJzcxUMBlVRURF9LBKJ6ODBg8rPz4/nSwEAujnPV8FdunRJtbW10Y/r6+t17Ngxpaena9CgQVq2bJl++tOf6sknn1Rubq7eeusthUIhTZs2LZ5zAwC6Oc8BOnToUMy9w0pLSyVJ8+bN08aNG/X666+rtbVVixYtUnNzs55//nnt2rWrQ/fzAgAkL88Bmjhxopxzd3ze5/Np9erVWr169QMNhs6TkZHRoXXHjh3zvObSpUue1/zpT3/yvObTTz/1vEaS2tvbO7Qu2Vy5csV6BDwEzK+CAwA8nAgQAMAEAQIAmCBAAAATBAgAYIIAAQBMECAAgAkCBAAwQYAAACYIEADABAECAJggQAAAEwQIAGDC5+52a2sDkUhEgUDAegygy/nzn//seU1RUVGHXis9Pd3zmubm5g69FpJXOBy+62+55gwIAGCCAAEATBAgAIAJAgQAMEGAAAAmCBAAwAQBAgCYIEAAABMECABgggABAEwQIACACQIEADDRy3oAAPcnOzvbegQgrjgDAgCYIEAAABMECABgggABAEwQIACACQIEADBBgAAAJggQAMAEAQIAmCBAAAATBAgAYIIAAQBMcDNSIIkdOnSoQ+taWlriPAlwK86AAAAmCBAAwAQBAgCYIEAAABMECABgggABAEwQIACACQIEADBBgAAAJggQAMAEAQIAmCBAAAAT3IwUMDBw4EDPa55++mnPaz799FPPayTpxo0bHVoHeMEZEADABAECAJjwHKB9+/ZpypQpCoVC8vl82rZtW8zz8+fPl8/ni9mKioriNS8AIEl4DlBra6tGjx6tdevW3XGfoqIinT17Nrpt3rz5gYYEACQfzxchFBcXq7i4+K77+P1+BYPBDg8FAEh+CfkeUGVlpTIzMzVs2DAtWbJEFy9evOO+bW1tikQiMRsAIPnFPUBFRUV6//33VVFRoZ///OeqqqpScXHxHS/rLC8vVyAQiG45OTnxHgkA0AXF/eeA5syZE/3zyJEjNWrUKA0dOlSVlZWaNGnSLfuXlZWptLQ0+nEkEiFCAPAQSPhl2EOGDFFGRoZqa2tv+7zf71daWlrMBgBIfgkP0OnTp3Xx4kVlZ2cn+qUAAN2I5y/BXbp0KeZspr6+XseOHVN6errS09O1atUqzZw5U8FgUHV1dXr99df1xBNPaPLkyXEdHADQvXkO0KFDh/Tiiy9GP/7q+zfz5s3T+vXrdfz4cf3+979Xc3OzQqGQCgsL9ZOf/ER+vz9+UwMAuj3PAZo4caKcc3d8/i9/+csDDQQ8DKZOnep5TUpKiuc17777ruc1QGfhXnAAABMECABgggABAEwQIACACQIEADBBgAAAJggQAMAEAQIAmCBAAAATBAgAYIIAAQBMECAAgAkCBAAwEfdfyQ3g3saPH+95TXt7u+c1DQ0NntcAnYUzIACACQIEADBBgAAAJggQAMAEAQIAmCBAAAATBAgAYIIAAQBMECAAgAkCBAAwQYAAACYIEADABDcjBQxkZ2d7XnP8+HHPa7gZKboyzoAAACYIEADABAECAJggQAAAEwQIAGCCAAEATBAgAIAJAgQAMEGAAAAmCBAAwAQBAgCYIEAAABMECABgggABAEwQIACACQIEADBBgAAAJggQAMAEAQIAmCBAAAATBAgAYIIAAQBMECAAgAkCBAAwQYAAACZ6WQ8AdHdpaWme1+Tl5Xles3//fs9rgK6MMyAAgAkCBAAw4SlA5eXlGjt2rFJTU5WZmalp06appqYmZp+rV6+qpKRE/fv312OPPaaZM2eqqakprkMDALo/TwGqqqpSSUmJDhw4oN27d+v69esqLCxUa2trdJ/ly5drx44d2rJli6qqqnTmzBnNmDEj7oMDALo3n3POdXTx+fPnlZmZqaqqKk2YMEHhcFgDBgzQpk2bNGvWLEnSF198oaefflrV1dV67rnn7vk5I5GIAoFAR0cCOl1HLkJobGz0vKYjFyEUFhZ6XgPESzgcvut/Hw/0PaBwOCxJSk9PlyQdPnxY169fV0FBQXSf4cOHa9CgQaqurr7t52hra1MkEonZAADJr8MBam9v17JlyzR+/HiNGDFC0s3/q0tJSVG/fv1i9s3Kyrrj//GVl5crEAhEt5ycnI6OBADoRjocoJKSEp04cUIfffTRAw1QVlamcDgc3U6dOvVAnw8A0D106AdRly5dqp07d2rfvn0aOHBg9PFgMKhr166pubk55iyoqalJwWDwtp/L7/fL7/d3ZAwAQDfm6QzIOaelS5dq69at2rt3r3Jzc2OeHzNmjHr37q2KioroYzU1NWpoaFB+fn58JgYAJAVPZ0AlJSXatGmTtm/frtTU1Oj3dQKBgPr27atAIKAFCxaotLRU6enpSktL0yuvvKL8/Pz7ugIOAPDw8BSg9evXS5ImTpwY8/iGDRs0f/58SdIvf/lL9ejRQzNnzlRbW5smT56sX//613EZFgCQPB7o54ASgZ8DQnczd+5cz2s++OADz2umTp3qec2OHTs8rwHiJaE/BwQAQEcRIACACQIEADBBgAAAJggQAMAEAQIAmCBAAAATBAgAYIIAAQBMECAAgAkCBAAwQYAAACYIEADARId+IyqA/5o1a1anvA6/rh7JhjMgAIAJAgQAMEGAAAAmCBAAwAQBAgCYIEAAABMECABgggABAEwQIACACQIEADBBgAAAJggQAMAENyMFDEQiEc9rzp8/n4BJADucAQEATBAgAIAJAgQAMEGAAAAmCBAAwAQBAgCYIEAAABMECABgggABAEwQIACACQIEADBBgAAAJrgZKfCAhg8f7nnNf/7zH89r/v3vf3teA3RlnAEBAEwQIACACQIEADBBgAAAJggQAMAEAQIAmCBAAAATBAgAYIIAAQBMECAAgAkCBAAwQYAAACa4GSnwP1599VXPazpyM9Lf/OY3ntcAyYYzIACACQIEADDhKUDl5eUaO3asUlNTlZmZqWnTpqmmpiZmn4kTJ8rn88VsixcvjuvQAIDuz1OAqqqqVFJSogMHDmj37t26fv26CgsL1draGrPfwoULdfbs2ei2Zs2auA4NAOj+PF2EsGvXrpiPN27cqMzMTB0+fFgTJkyIPv7II48oGAzGZ0IAQFJ6oO8BhcNhSVJ6enrM4x9++KEyMjI0YsQIlZWV6fLly3f8HG1tbYpEIjEbACD5dfgy7Pb2di1btkzjx4/XiBEjoo+//PLLGjx4sEKhkI4fP6433nhDNTU1+uSTT277ecrLy7Vq1aqOjgEA6KY6HKCSkhKdOHFC+/fvj3l80aJF0T+PHDlS2dnZmjRpkurq6jR06NBbPk9ZWZlKS0ujH0ciEeXk5HR0LABAN9GhAC1dulQ7d+7Uvn37NHDgwLvum5eXJ0mqra29bYD8fr/8fn9HxgAAdGOeAuSc0yuvvKKtW7eqsrJSubm591xz7NgxSVJ2dnaHBgQAJCdPASopKdGmTZu0fft2paamqrGxUZIUCATUt29f1dXVadOmTfrud7+r/v376/jx41q+fLkmTJigUaNGJeQfAADQPXkK0Pr16yXd/GHT/7VhwwbNnz9fKSkp2rNnj9auXavW1lbl5ORo5syZevPNN+M2MAAgOXj+Etzd5OTkqKqq6oEGAgA8HLgbNvA/+vfv3ymv88c//rFTXgfoyrgZKQDABAECAJggQAAAEwQIAGCCAAEATBAgAIAJAgQAMEGAAAAmCBAAwAQBAgCYIEAAABMECABgwufudYvrThaJRBQIBKzHAAA8oHA4rLS0tDs+zxkQAMAEAQIAmCBAAAATBAgAYIIAAQBMECAAgAkCBAAwQYAAACYIEADABAECAJggQAAAE10uQF3s1nQAgA6619/nXS5ALS0t1iMAAOLgXn+fd7m7Ybe3t+vMmTNKTU2Vz+eLeS4SiSgnJ0enTp266x1Wkx3H4SaOw00ch5s4Djd1hePgnFNLS4tCoZB69LjzeU6vTpzpvvTo0UMDBw686z5paWkP9RvsKxyHmzgON3EcbuI43GR9HO7n1+p0uS/BAQAeDgQIAGCiWwXI7/dr5cqV8vv91qOY4jjcxHG4ieNwE8fhpu50HLrcRQgAgIdDtzoDAgAkDwIEADBBgAAAJggQAMBEtwnQunXr9Pjjj6tPnz7Ky8vT559/bj1Sp3v77bfl8/lituHDh1uPlXD79u3TlClTFAqF5PP5tG3btpjnnXNasWKFsrOz1bdvXxUUFOjkyZM2wybQvY7D/Pnzb3l/FBUV2QybIOXl5Ro7dqxSU1OVmZmpadOmqaamJmafq1evqqSkRP3799djjz2mmTNnqqmpyWjixLif4zBx4sRb3g+LFy82mvj2ukWAPv74Y5WWlmrlypU6cuSIRo8ercmTJ+vcuXPWo3W6Z555RmfPno1u+/fvtx4p4VpbWzV69GitW7futs+vWbNG7777rt577z0dPHhQjz76qCZPnqyrV6928qSJda/jIElFRUUx74/Nmzd34oSJV1VVpZKSEh04cEC7d+/W9evXVVhYqNbW1ug+y5cv144dO7RlyxZVVVXpzJkzmjFjhuHU8Xc/x0GSFi5cGPN+WLNmjdHEd+C6gXHjxrmSkpLoxzdu3HChUMiVl5cbTtX5Vq5c6UaPHm09hilJbuvWrdGP29vbXTAYdL/4xS+ijzU3Nzu/3+82b95sMGHn+PpxcM65efPmualTp5rMY+XcuXNOkquqqnLO3fx337t3b7dly5boPv/85z+dJFddXW01ZsJ9/Tg459x3vvMd98Mf/tBuqPvQ5c+Arl27psOHD6ugoCD6WI8ePVRQUKDq6mrDyWycPHlSoVBIQ4YM0dy5c9XQ0GA9kqn6+no1NjbGvD8CgYDy8vIeyvdHZWWlMjMzNWzYMC1ZskQXL160HimhwuGwJCk9PV2SdPjwYV2/fj3m/TB8+HANGjQoqd8PXz8OX/nwww+VkZGhESNGqKysTJcvX7YY74663M1Iv+7ChQu6ceOGsrKyYh7PysrSF198YTSVjby8PG3cuFHDhg3T2bNntWrVKr3wwgs6ceKEUlNTrccz0djYKEm3fX989dzDoqioSDNmzFBubq7q6ur04x//WMXFxaqurlbPnj2tx4u79vZ2LVu2TOPHj9eIESMk3Xw/pKSkqF+/fjH7JvP74XbHQZJefvllDR48WKFQSMePH9cbb7yhmpoaffLJJ4bTxuryAcJ/FRcXR/88atQo5eXlafDgwfrDH/6gBQsWGE6GrmDOnDnRP48cOVKjRo3S0KFDVVlZqUmTJhlOlhglJSU6ceLEQ/F90Lu503FYtGhR9M8jR45Udna2Jk2apLq6Og0dOrSzx7ytLv8luIyMDPXs2fOWq1iampoUDAaNpuoa+vXrp6eeekq1tbXWo5j56j3A++NWQ4YMUUZGRlK+P5YuXaqdO3fqs88+i/n1LcFgUNeuXVNzc3PM/sn6frjTcbidvLw8SepS74cuH6CUlBSNGTNGFRUV0cfa29tVUVGh/Px8w8nsXbp0SXV1dcrOzrYexUxubq6CwWDM+yMSiejgwYMP/fvj9OnTunjxYlK9P5xzWrp0qbZu3aq9e/cqNzc35vkxY8aod+/eMe+HmpoaNTQ0JNX74V7H4XaOHTsmSV3r/WB9FcT9+Oijj5zf73cbN250//jHP9yiRYtcv379XGNjo/VonerVV191lZWVrr6+3v31r391BQUFLiMjw507d856tIRqaWlxR48edUePHnWS3DvvvOOOHj3q/vWvfznnnPvZz37m+vXr57Zv3+6OHz/upk6d6nJzc92VK1eMJ4+vux2HlpYW99prr7nq6mpXX1/v9uzZ4771rW+5J5980l29etV69LhZsmSJCwQCrrKy0p09eza6Xb58ObrP4sWL3aBBg9zevXvdoUOHXH5+vsvPzzecOv7udRxqa2vd6tWr3aFDh1x9fb3bvn27GzJkiJswYYLx5LG6RYCcc+5Xv/qVGzRokEtJSXHjxo1zBw4csB6p082ePdtlZ2e7lJQU981vftPNnj3b1dbWWo+VcJ999pmTdMs2b94859zNS7Hfeustl5WV5fx+v5s0aZKrqamxHToB7nYcLl++7AoLC92AAQNc79693eDBg93ChQuT7n/SbvfPL8lt2LAhus+VK1fcD37wA/eNb3zDPfLII2769Onu7NmzdkMnwL2OQ0NDg5swYYJLT093fr/fPfHEE+5HP/qRC4fDtoN/Db+OAQBgost/DwgAkJwIEADABAECAJggQAAAEwQIAGCCAAEATBAgAIAJAgQAMEGAAAAmCBAAwAQBAgCYIEAAABP/DyNEW274zDV4AAAAAElFTkSuQmCC",
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "%matplotlib inline\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "\n",
        "print(f'training inputs: shape {x_train.shape}, dtype {x_train.dtype}')\n",
        "print(f'training outputs: shape {y_train.shape}, dtype {y_train.dtype}')\n",
        "print(f'input range: {x_train.min(), x_train.max()}')\n",
        "print(f'label range: {y_train.min(), y_train.max()}')\n",
        "\n",
        "i_example = 42\n",
        "x_train_example = x_train[i_example].reshape((28, 28))\n",
        "y_train_example = y_train[i_example]\n",
        "plt.imshow(x_train_example, cmap='gray')\n",
        "print(f'label of training example {i_example}: {y_train_example}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "COOGtiMZXqYq"
      },
      "source": [
        "We need to convert the data into PyTorch tensors. For ease of use, tensors have some familiar attributes and methods."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VEk34MTIVfh2",
        "outputId": "517ba60d-922d-4185-fa97-4787689bb2ed"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "training inputs: shape torch.Size([50000, 784]), dtype torch.float32\n",
            "training outputs: shape torch.Size([50000]), dtype torch.int64\n",
            "input range: (tensor(0.), tensor(0.9961))\n",
            "label range: (tensor(0), tensor(9))\n"
          ]
        }
      ],
      "source": [
        "x_train, y_train, x_val, y_val = map(torch.tensor, (x_train, y_train, x_val, y_val))\n",
        "print(f'training inputs: shape {x_train.shape}, dtype {x_train.dtype}')\n",
        "print(f'training outputs: shape {y_train.shape}, dtype {y_train.dtype}')\n",
        "print(f'input range: {x_train.min(), x_train.max()}')\n",
        "print(f'label range: {y_train.min(), y_train.max()}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YuLs5f2bcwXU"
      },
      "source": [
        "### Bare Metal MNIST\n",
        "We will first train a model for classifying MNIST digits with nothing but tensor and gradient operations. We'll later substitute out chunks of this code with the appropriate PyTorch abstractions.\n",
        "\n",
        "To keep things simple, for this we'll use a linear model. For a single example, the model is\n",
        "$$ z = W^\\top x + b $$\n",
        "where $W \\in \\mathbb{R}^{784 \\times 10}$ and $b \\in \\mathbb{R}^{10}$ are the parameters to be learned, $x \\in [0, 1]^{784}$ is the MNIST image, and $z \\in \\mathbb{R}^{10}$ are the logits of the categorical distribution for the label of $x$.\n",
        "\n",
        "To obtain categorical parameters from the logits, we do\n",
        "$$ p = \\text{softmax}(z) = \\frac{1}{\\sum_j \\exp(z_j)} \\exp(z).$$\n",
        "\n",
        "The label $y \\in \\{0, \\dots, 9\\}$ implicitly defines a one-hot categorical parameter vector $y' \\in \\{0, 1\\}^{10}$, where $y$ corresponds to the index of $y'$ that is 1. The negative log-likelihood is then\n",
        "$$ \\text{nll}(p, y) = - \\sum_{i=0}^{9} y'_i \\log(p_i) = -\\log(p_y)$$\n",
        "\n",
        "Note that our implementation will assume batched inputs with a batch size of $B$. That is,\n",
        "$$x_b = \\begin{bmatrix} x_1^\\top \\\\ \\vdots \\\\ x_B^\\top \\end{bmatrix} \\in [0, 1]^{B \\times 784}.$$ We can vectorize the above operations to process all $B$ inputs at once (naively using `for` loops is much, much slower). For example, to compute the logits, we have\n",
        "$$z_b = x_b W + \\begin{bmatrix} b^\\top \\\\ \\vdots \\\\ b^\\top \\end{bmatrix} \\in \\mathbb{R}^{B \\times 10}$$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "dZEHIaizc3CU"
      },
      "outputs": [],
      "source": [
        "W = torch.randn(784, 10, requires_grad=True)\n",
        "b = torch.randn(10, requires_grad=True)\n",
        "\n",
        "def logits(x_b, W, b):\n",
        "    return x_b @ W + b    # addition is broadcasted\n",
        "\n",
        "def log_softmax(z_b):\n",
        "    return z_b - torch.logsumexp(z_b, dim=1, keepdim=True)  # combining log and softmax is more numerically stable\n",
        "\n",
        "def negative_log_likelihood(logp_b, y_b):\n",
        "    return -torch.mean(logp_b[range(y_b.shape[0]), y_b])   # indexing trick\n",
        "\n",
        "def accuracy(logit_b, y_b):\n",
        "    return (logit_b.argmax(dim=1) == y_b).float().mean().item()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zQYMuIF7i87C"
      },
      "source": [
        "Note that the `requires_grad` keyword argument/attribute determines whether a tensor is a constant or a variable."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0vC_MuGOjDTn",
        "outputId": "2ce5407b-bf8e-4789-d930-4707fa2f8a8c"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "True\n",
            "False\n"
          ]
        }
      ],
      "source": [
        "print(b.requires_grad)\n",
        "print(y_train.requires_grad)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DrS8cR3OiARG"
      },
      "source": [
        "### Computation Graph\n",
        "Now we have the opportunity to look at some of PyTorch's core mechanics: reverse-mode automatic differentiation (a.k.a. backpropagation) on a dynamic computation graph. Let's take a batch of training data and compute the average loss over the batch."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "LGQELJgRg_TY"
      },
      "outputs": [],
      "source": [
        "batch_size = 64\n",
        "x_b, y_b = x_train[:batch_size], y_train[:batch_size]\n",
        "\n",
        "def loss_fn(W, b, x_b, y_b):\n",
        "    logit_b = logits(x_b, W, b)\n",
        "    logp_b = log_softmax(logit_b)\n",
        "    loss = negative_log_likelihood(logp_b, y_b)\n",
        "    return loss\n",
        "\n",
        "loss = loss_fn(W, b, x_b, y_b)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7YpM1c2sj-GN"
      },
      "source": [
        "Computations like the above are automagically recorded on a computation graph. Let's take a peek."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Iab2ejfRlpLv",
        "outputId": "edaebecb-249d-4d6f-8d8b-ebf69e8f73c4"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: torchviz in /usr/local/lib/python3.11/dist-packages (0.0.3)\n",
            "Requirement already satisfied: torch in /usr/local/lib/python3.11/dist-packages (from torchviz) (2.5.1+cu124)\n",
            "Requirement already satisfied: graphviz in /usr/local/lib/python3.11/dist-packages (from torchviz) (0.20.3)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.11/dist-packages (from torch->torchviz) (3.17.0)\n",
            "Requirement already satisfied: typing-extensions>=4.8.0 in /usr/local/lib/python3.11/dist-packages (from torch->torchviz) (4.12.2)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.11/dist-packages (from torch->torchviz) (3.4.2)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.11/dist-packages (from torch->torchviz) (3.1.5)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.11/dist-packages (from torch->torchviz) (2024.10.0)\n",
            "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.4.127 in /usr/local/lib/python3.11/dist-packages (from torch->torchviz) (12.4.127)\n",
            "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.4.127 in /usr/local/lib/python3.11/dist-packages (from torch->torchviz) (12.4.127)\n",
            "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.4.127 in /usr/local/lib/python3.11/dist-packages (from torch->torchviz) (12.4.127)\n",
            "Requirement already satisfied: nvidia-cudnn-cu12==9.1.0.70 in /usr/local/lib/python3.11/dist-packages (from torch->torchviz) (9.1.0.70)\n",
            "Requirement already satisfied: nvidia-cublas-cu12==12.4.5.8 in /usr/local/lib/python3.11/dist-packages (from torch->torchviz) (12.4.5.8)\n",
            "Requirement already satisfied: nvidia-cufft-cu12==11.2.1.3 in /usr/local/lib/python3.11/dist-packages (from torch->torchviz) (11.2.1.3)\n",
            "Requirement already satisfied: nvidia-curand-cu12==10.3.5.147 in /usr/local/lib/python3.11/dist-packages (from torch->torchviz) (10.3.5.147)\n",
            "Requirement already satisfied: nvidia-cusolver-cu12==11.6.1.9 in /usr/local/lib/python3.11/dist-packages (from torch->torchviz) (11.6.1.9)\n",
            "Requirement already satisfied: nvidia-cusparse-cu12==12.3.1.170 in /usr/local/lib/python3.11/dist-packages (from torch->torchviz) (12.3.1.170)\n",
            "Requirement already satisfied: nvidia-nccl-cu12==2.21.5 in /usr/local/lib/python3.11/dist-packages (from torch->torchviz) (2.21.5)\n",
            "Requirement already satisfied: nvidia-nvtx-cu12==12.4.127 in /usr/local/lib/python3.11/dist-packages (from torch->torchviz) (12.4.127)\n",
            "Requirement already satisfied: nvidia-nvjitlink-cu12==12.4.127 in /usr/local/lib/python3.11/dist-packages (from torch->torchviz) (12.4.127)\n",
            "Requirement already satisfied: triton==3.1.0 in /usr/local/lib/python3.11/dist-packages (from torch->torchviz) (3.1.0)\n",
            "Requirement already satisfied: sympy==1.13.1 in /usr/local/lib/python3.11/dist-packages (from torch->torchviz) (1.13.1)\n",
            "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.11/dist-packages (from sympy==1.13.1->torch->torchviz) (1.3.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.11/dist-packages (from jinja2->torch->torchviz) (3.0.2)\n"
          ]
        }
      ],
      "source": [
        "!pip install torchviz"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 765
        },
        "id": "5zVVH_dkkaeq",
        "outputId": "174639b5-2c5c-46a4-f72b-1c291ce80fd3"
      },
      "outputs": [
        {
          "data": {
            "image/svg+xml": [
              "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
              "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
              " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
              "<!-- Generated by graphviz version 2.43.0 (0)\n",
              " -->\n",
              "<!-- Title: %3 Pages: 1 -->\n",
              "<svg width=\"222pt\" height=\"558pt\"\n",
              " viewBox=\"0.00 0.00 222.00 558.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
              "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 554)\">\n",
              "<title>%3</title>\n",
              "<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-554 218,-554 218,4 -4,4\"/>\n",
              "<!-- 134762862158448 -->\n",
              "<g id=\"node1\" class=\"node\">\n",
              "<title>134762862158448</title>\n",
              "<polygon fill=\"#caff70\" stroke=\"black\" points=\"132.5,-31 78.5,-31 78.5,0 132.5,0 132.5,-31\"/>\n",
              "<text text-anchor=\"middle\" x=\"105.5\" y=\"-7\" font-family=\"monospace\" font-size=\"10.00\"> ()</text>\n",
              "</g>\n",
              "<!-- 134762670192480 -->\n",
              "<g id=\"node2\" class=\"node\">\n",
              "<title>134762670192480</title>\n",
              "<polygon fill=\"lightgrey\" stroke=\"black\" points=\"150,-86 61,-86 61,-67 150,-67 150,-86\"/>\n",
              "<text text-anchor=\"middle\" x=\"105.5\" y=\"-74\" font-family=\"monospace\" font-size=\"10.00\">NegBackward0</text>\n",
              "</g>\n",
              "<!-- 134762670192480&#45;&gt;134762862158448 -->\n",
              "<g id=\"edge12\" class=\"edge\">\n",
              "<title>134762670192480&#45;&gt;134762862158448</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M105.5,-66.79C105.5,-60.07 105.5,-50.4 105.5,-41.34\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"109,-41.19 105.5,-31.19 102,-41.19 109,-41.19\"/>\n",
              "</g>\n",
              "<!-- 134762670192336 -->\n",
              "<g id=\"node3\" class=\"node\">\n",
              "<title>134762670192336</title>\n",
              "<polygon fill=\"lightgrey\" stroke=\"black\" points=\"153,-141 58,-141 58,-122 153,-122 153,-141\"/>\n",
              "<text text-anchor=\"middle\" x=\"105.5\" y=\"-129\" font-family=\"monospace\" font-size=\"10.00\">MeanBackward0</text>\n",
              "</g>\n",
              "<!-- 134762670192336&#45;&gt;134762670192480 -->\n",
              "<g id=\"edge1\" class=\"edge\">\n",
              "<title>134762670192336&#45;&gt;134762670192480</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M105.5,-121.75C105.5,-114.8 105.5,-104.85 105.5,-96.13\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"109,-96.09 105.5,-86.09 102,-96.09 109,-96.09\"/>\n",
              "</g>\n",
              "<!-- 134762670192192 -->\n",
              "<g id=\"node4\" class=\"node\">\n",
              "<title>134762670192192</title>\n",
              "<polygon fill=\"lightgrey\" stroke=\"black\" points=\"156,-196 55,-196 55,-177 156,-177 156,-196\"/>\n",
              "<text text-anchor=\"middle\" x=\"105.5\" y=\"-184\" font-family=\"monospace\" font-size=\"10.00\">IndexBackward0</text>\n",
              "</g>\n",
              "<!-- 134762670192192&#45;&gt;134762670192336 -->\n",
              "<g id=\"edge2\" class=\"edge\">\n",
              "<title>134762670192192&#45;&gt;134762670192336</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M105.5,-176.75C105.5,-169.8 105.5,-159.85 105.5,-151.13\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"109,-151.09 105.5,-141.09 102,-151.09 109,-151.09\"/>\n",
              "</g>\n",
              "<!-- 134762670192144 -->\n",
              "<g id=\"node5\" class=\"node\">\n",
              "<title>134762670192144</title>\n",
              "<polygon fill=\"lightgrey\" stroke=\"black\" points=\"150,-251 61,-251 61,-232 150,-232 150,-251\"/>\n",
              "<text text-anchor=\"middle\" x=\"105.5\" y=\"-239\" font-family=\"monospace\" font-size=\"10.00\">SubBackward0</text>\n",
              "</g>\n",
              "<!-- 134762670192144&#45;&gt;134762670192192 -->\n",
              "<g id=\"edge3\" class=\"edge\">\n",
              "<title>134762670192144&#45;&gt;134762670192192</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M105.5,-231.75C105.5,-224.8 105.5,-214.85 105.5,-206.13\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"109,-206.09 105.5,-196.09 102,-206.09 109,-206.09\"/>\n",
              "</g>\n",
              "<!-- 134762670192048 -->\n",
              "<g id=\"node6\" class=\"node\">\n",
              "<title>134762670192048</title>\n",
              "<polygon fill=\"lightgrey\" stroke=\"black\" points=\"150,-361 61,-361 61,-342 150,-342 150,-361\"/>\n",
              "<text text-anchor=\"middle\" x=\"105.5\" y=\"-349\" font-family=\"monospace\" font-size=\"10.00\">AddBackward0</text>\n",
              "</g>\n",
              "<!-- 134762670192048&#45;&gt;134762670192144 -->\n",
              "<g id=\"edge4\" class=\"edge\">\n",
              "<title>134762670192048&#45;&gt;134762670192144</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M98.16,-341.68C89.02,-329.62 75.22,-307.17 80.5,-287 83.05,-277.26 88.25,-267.41 93.26,-259.46\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"96.19,-261.37 98.87,-251.12 90.39,-257.46 96.19,-261.37\"/>\n",
              "</g>\n",
              "<!-- 134762670192096 -->\n",
              "<g id=\"node12\" class=\"node\">\n",
              "<title>134762670192096</title>\n",
              "<polygon fill=\"lightgrey\" stroke=\"black\" points=\"214,-306 89,-306 89,-287 214,-287 214,-306\"/>\n",
              "<text text-anchor=\"middle\" x=\"151.5\" y=\"-294\" font-family=\"monospace\" font-size=\"10.00\">LogsumexpBackward0</text>\n",
              "</g>\n",
              "<!-- 134762670192048&#45;&gt;134762670192096 -->\n",
              "<g id=\"edge11\" class=\"edge\">\n",
              "<title>134762670192048&#45;&gt;134762670192096</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M113.1,-341.75C119.66,-334.18 129.32,-323.05 137.31,-313.84\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"140.13,-315.94 144.05,-306.09 134.85,-311.35 140.13,-315.94\"/>\n",
              "</g>\n",
              "<!-- 134762670191760 -->\n",
              "<g id=\"node7\" class=\"node\">\n",
              "<title>134762670191760</title>\n",
              "<polygon fill=\"lightgrey\" stroke=\"black\" points=\"92,-416 9,-416 9,-397 92,-397 92,-416\"/>\n",
              "<text text-anchor=\"middle\" x=\"50.5\" y=\"-404\" font-family=\"monospace\" font-size=\"10.00\">MmBackward0</text>\n",
              "</g>\n",
              "<!-- 134762670191760&#45;&gt;134762670192048 -->\n",
              "<g id=\"edge5\" class=\"edge\">\n",
              "<title>134762670191760&#45;&gt;134762670192048</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M59.58,-396.75C67.59,-389.03 79.46,-377.6 89.12,-368.28\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"91.81,-370.55 96.59,-361.09 86.96,-365.51 91.81,-370.55\"/>\n",
              "</g>\n",
              "<!-- 134762670191616 -->\n",
              "<g id=\"node8\" class=\"node\">\n",
              "<title>134762670191616</title>\n",
              "<polygon fill=\"lightgrey\" stroke=\"black\" points=\"101,-477 0,-477 0,-458 101,-458 101,-477\"/>\n",
              "<text text-anchor=\"middle\" x=\"50.5\" y=\"-465\" font-family=\"monospace\" font-size=\"10.00\">AccumulateGrad</text>\n",
              "</g>\n",
              "<!-- 134762670191616&#45;&gt;134762670191760 -->\n",
              "<g id=\"edge6\" class=\"edge\">\n",
              "<title>134762670191616&#45;&gt;134762670191760</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M50.5,-457.79C50.5,-449.6 50.5,-437.06 50.5,-426.55\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"54,-426.24 50.5,-416.24 47,-426.24 54,-426.24\"/>\n",
              "</g>\n",
              "<!-- 134762861964240 -->\n",
              "<g id=\"node9\" class=\"node\">\n",
              "<title>134762861964240</title>\n",
              "<polygon fill=\"lightblue\" stroke=\"black\" points=\"89,-550 12,-550 12,-519 89,-519 89,-550\"/>\n",
              "<text text-anchor=\"middle\" x=\"50.5\" y=\"-526\" font-family=\"monospace\" font-size=\"10.00\"> (784, 10)</text>\n",
              "</g>\n",
              "<!-- 134762861964240&#45;&gt;134762670191616 -->\n",
              "<g id=\"edge7\" class=\"edge\">\n",
              "<title>134762861964240&#45;&gt;134762670191616</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M50.5,-518.75C50.5,-509.39 50.5,-497.19 50.5,-487.16\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"54,-487.02 50.5,-477.02 47,-487.02 54,-487.02\"/>\n",
              "</g>\n",
              "<!-- 134762670191808 -->\n",
              "<g id=\"node10\" class=\"node\">\n",
              "<title>134762670191808</title>\n",
              "<polygon fill=\"lightgrey\" stroke=\"black\" points=\"211,-416 110,-416 110,-397 211,-397 211,-416\"/>\n",
              "<text text-anchor=\"middle\" x=\"160.5\" y=\"-404\" font-family=\"monospace\" font-size=\"10.00\">AccumulateGrad</text>\n",
              "</g>\n",
              "<!-- 134762670191808&#45;&gt;134762670192048 -->\n",
              "<g id=\"edge8\" class=\"edge\">\n",
              "<title>134762670191808&#45;&gt;134762670192048</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M151.42,-396.75C143.41,-389.03 131.54,-377.6 121.88,-368.28\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"124.04,-365.51 114.41,-361.09 119.19,-370.55 124.04,-365.51\"/>\n",
              "</g>\n",
              "<!-- 134762862158736 -->\n",
              "<g id=\"node11\" class=\"node\">\n",
              "<title>134762862158736</title>\n",
              "<polygon fill=\"lightblue\" stroke=\"black\" points=\"187.5,-483 133.5,-483 133.5,-452 187.5,-452 187.5,-483\"/>\n",
              "<text text-anchor=\"middle\" x=\"160.5\" y=\"-459\" font-family=\"monospace\" font-size=\"10.00\"> (10)</text>\n",
              "</g>\n",
              "<!-- 134762862158736&#45;&gt;134762670191808 -->\n",
              "<g id=\"edge9\" class=\"edge\">\n",
              "<title>134762862158736&#45;&gt;134762670191808</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M160.5,-451.92C160.5,-444.22 160.5,-434.69 160.5,-426.43\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"164,-426.25 160.5,-416.25 157,-426.25 164,-426.25\"/>\n",
              "</g>\n",
              "<!-- 134762670192096&#45;&gt;134762670192144 -->\n",
              "<g id=\"edge10\" class=\"edge\">\n",
              "<title>134762670192096&#45;&gt;134762670192144</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M143.9,-286.75C137.34,-279.18 127.68,-268.05 119.69,-258.84\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"122.15,-256.35 112.95,-251.09 116.87,-260.94 122.15,-256.35\"/>\n",
              "</g>\n",
              "</g>\n",
              "</svg>\n"
            ],
            "text/plain": [
              "<graphviz.graphs.Digraph at 0x7a90ec80f5d0>"
            ]
          },
          "execution_count": 10,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "import torchviz\n",
        "torchviz.make_dot(loss)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BwJ-5hG7mTOL"
      },
      "source": [
        "We can actually inspect tensors to see that each non-leaf variable tensor records its gradient function based on how it was obtained."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gbRLFWkWmUY8",
        "outputId": "1d8fdd11-9ced-42f1-9391-eb8dd9590645"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "None\n",
            "<NegBackward0 object at 0x7a90e115f760>\n",
            "None\n"
          ]
        }
      ],
      "source": [
        "print(b.grad_fn)    # leaf variable\n",
        "print(loss.grad_fn)\n",
        "print(x_b.grad_fn)  # not a variable"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0u88oEuRSZAG"
      },
      "source": [
        "### Using `autograd.grad` to Compute Gradients\n",
        "We can now use reverse-mode automatic differentiation to compute the gradient of the loss with respect to our parameters!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "a4IuGz-vNru5",
        "outputId": "1fae385a-c3b8-4cce-ba33-f6a2c563d69e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor([ 0.2051, -0.1250,  0.1485, -0.0757, -0.1249,  0.0024, -0.0405,  0.0582,\n",
            "        -0.0259, -0.0222])\n"
          ]
        }
      ],
      "source": [
        "from torch import autograd\n",
        "loss = loss_fn(W, b, x_b, y_b)\n",
        "W_grad, b_grad = autograd.grad(loss, inputs=(W, b))\n",
        "print(b_grad)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I1U8J8ujfoSp"
      },
      "source": [
        "We're now ready to optimize our linear MNIST model. Each iteration consists of the following steps:\n",
        "* Sample a batch of training data.\n",
        "* Compute the logits from the inputs using the model parameters.\n",
        "* Compute a scalar loss from the logits and labels.\n",
        "* Compute the gradient of the loss with respect to the parameters.\n",
        "* Update the model parameters using the gradient.\n",
        "\n",
        "Note that the last step is done in a context manager that disables the construction of the computation graph.\n",
        "\n",
        "To check whether our optimization was successful, we'll look at the accuracy on the training data before and after."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iABo_4DPoInh",
        "outputId": "c0523790-e0b4-4a16-91bf-d905df6a3f0f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "accuracy before: 0.10620000213384628\n",
            "accuracy after: 0.8812599778175354\n"
          ]
        }
      ],
      "source": [
        "lr = 0.5\n",
        "num_epochs = 2\n",
        "\n",
        "def train(W, b):\n",
        "    print(f'accuracy before: {accuracy(logits(x_train, W, b), y_train)}')\n",
        "\n",
        "    for i_epoch in range(num_epochs):\n",
        "        i_batch_start = 0\n",
        "        while i_batch_start + batch_size < x_train.shape[0]:\n",
        "            x_b = x_train[i_batch_start:i_batch_start + batch_size]\n",
        "            y_b = y_train[i_batch_start:i_batch_start + batch_size]\n",
        "            i_batch_start += batch_size\n",
        "\n",
        "            logit_b = logits(x_b, W, b)\n",
        "            logp_b = log_softmax(logit_b)\n",
        "            loss = negative_log_likelihood(logp_b, y_b)\n",
        "            W_grad, b_grad = autograd.grad(loss, inputs=(W, b))\n",
        "\n",
        "            with torch.no_grad():   # we don't need gradients for this\n",
        "                W -= lr * W_grad\n",
        "                b -= lr * b_grad\n",
        "\n",
        "    print(f'accuracy after: {accuracy(logits(x_train, W, b), y_train)}')\n",
        "\n",
        "W = torch.randn(784, 10, requires_grad=True)\n",
        "b = torch.randn(10, requires_grad=True)\n",
        "train(W, b)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sYArnyJ_nDCt"
      },
      "source": [
        "### Using `backward` to Compute Gradients\n",
        "An alternative to using `autograd.grad` is to call `backward()` on the loss tensor. This _accumulates_ gradients of leaf tensors into their `grad` attribute. (Note that with this we don't need to specify the variables we're taking the gradient with respect to, nor do we obtain the gradients as output from the function call.)\n",
        "\n",
        "When we're done using a gradient in a `grad` attribute, we should reset it to 0, else repeated accumulations may occur. We can do this with the `Tensor` method `zero_()`. (Note that the `_` signifies that the operation modifies the tensor in-place.)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ta5qt6V1kPAE",
        "outputId": "4a766398-20de-44a4-e223-1f579cd8acf1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "before computing gradients\n",
            "None\n",
            "after forward pass and calling backward()\n",
            "tensor([-0.0010, -0.0039,  0.0037,  0.0023, -0.0114,  0.0132, -0.0002,  0.0235,\n",
            "         0.0132, -0.0394])\n",
            "after another forward pass and calling backward() again\n",
            "tensor([-0.0019, -0.0078,  0.0073,  0.0046, -0.0227,  0.0263, -0.0005,  0.0470,\n",
            "         0.0265, -0.0787])\n",
            "after calling zero_()\n",
            "tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])\n"
          ]
        }
      ],
      "source": [
        "print('before computing gradients')\n",
        "print(b.grad)\n",
        "\n",
        "loss = loss_fn(W, b, x_b, y_b)\n",
        "loss.backward()\n",
        "print('after forward pass and calling backward()')\n",
        "print(b.grad)\n",
        "\n",
        "loss = loss_fn(W, b, x_b, y_b)\n",
        "loss.backward()\n",
        "print('after another forward pass and calling backward() again')\n",
        "print(b.grad)\n",
        "\n",
        "b.grad.zero_()\n",
        "print('after calling zero_()')\n",
        "print(b.grad)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I9gxa-qsSc0i"
      },
      "source": [
        "When should you use `autograd.grad` vs. `backward`? In general, `backward` is more convenient, but there are times (e.g. in Homework 2) when you need more fine-grained control over gradient computation. Let's see what the training code looks like when we use `backward`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iYhQfdJDLJyj",
        "outputId": "002938bc-52a7-45ad-e9ea-4b3bbacbcdb1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "accuracy before: 0.07010000199079514\n",
            "accuracy after: 0.8821200132369995\n"
          ]
        }
      ],
      "source": [
        "def train(W, b):\n",
        "    print(f'accuracy before: {accuracy(logits(x_train, W, b), y_train)}')\n",
        "\n",
        "    for i_epoch in range(num_epochs):\n",
        "        i_batch_start = 0\n",
        "        while i_batch_start + batch_size < x_train.shape[0]:\n",
        "            x_b = x_train[i_batch_start:i_batch_start + batch_size]\n",
        "            y_b = y_train[i_batch_start:i_batch_start + batch_size]\n",
        "            i_batch_start += batch_size\n",
        "\n",
        "            logit_b = logits(x_b, W, b)\n",
        "            logp_b = log_softmax(logit_b)\n",
        "            loss = negative_log_likelihood(logp_b, y_b)\n",
        "            loss.backward()\n",
        "\n",
        "            with torch.no_grad():\n",
        "                W -= lr * W.grad\n",
        "                b -= lr * b.grad\n",
        "                W.grad.zero_()\n",
        "                b.grad.zero_()\n",
        "\n",
        "    print(f'accuracy after: {accuracy(logits(x_train, W, b), y_train)}')\n",
        "\n",
        "W = torch.randn(784, 10, requires_grad=True)\n",
        "b = torch.randn(10, requires_grad=True)\n",
        "train(W, b)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v1MUJ0cfxD58"
      },
      "source": [
        "### `nn.functional`\n",
        "Now, let's refactor and extend our training code by allowing ourselves to use other `torch` modules. We'll first look at `torch.nn`.\n",
        "\n",
        "`nn.functional` (commonly abbreviated as just `F`) includes many useful functions, e.g. loss functions and activation functions. In particular, `F.cross_entropy` combines the `log_softmax` and `negative_log_likelihood` functions into one.\n",
        "\n",
        "If you're used to using TensorFlow's [CategoricalCrossentropy loss](https://www.tensorflow.org/api_docs/python/tf/keras/losses/CategoricalCrossentropy), note that `F.cross_entropy` takes integer labels, not one-hot labels"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5A2M3dY4xpfs",
        "outputId": "aaabdfb3-9419-4e3b-d400-05ae899c4ba5"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor(0.1862)\n",
            "tensor(0.1862)\n"
          ]
        }
      ],
      "source": [
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "with torch.no_grad():\n",
        "    logit_b = logits(x_b, W, b)\n",
        "    print(negative_log_likelihood(log_softmax(logit_b), y_b))\n",
        "    print(F.cross_entropy(logit_b, y_b))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I91UsJOmyv1b"
      },
      "source": [
        "### `nn.Module`\n",
        "`nn.Module` is a useful class that corresponds loosely to an intuitive notion of a parameterized model. It\n",
        "* registers `nn.Parameter`s, which are essentially variable tensors, as attributes\n",
        "* holds as the `forward` method the model's forward pass"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "MIN1V4b7yvGF"
      },
      "outputs": [],
      "source": [
        "class MNISTLinear(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.W = nn.Parameter(torch.randn(784, 10))\n",
        "        self.b = nn.Parameter(torch.randn(10))\n",
        "\n",
        "    def forward(self, x_b):\n",
        "        return x_b @ self.W + self.b"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mpd724uU3Vkg"
      },
      "source": [
        "Using the `nn.Module` can be very convenient. For example, previously we had to manually zero the `grad` attribute of `W` and `b`. We can now use the `parameters` method, which returns an iterator over module parameters, and the `zero_grad` method, which zeros the `grad` attribute of every module parameter.\n",
        "\n",
        "After substituting code for the model forward pass, the loss computation, the optimization step, and the gradient zeroing, the training code now looks like this:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5SJlK8y24BHn",
        "outputId": "35634d90-27f3-4203-964e-59e8bdeb1dd4"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "accuracy before: 0.07667999714612961\n",
            "accuracy after: 0.8832799792289734\n"
          ]
        }
      ],
      "source": [
        "def train(model):\n",
        "    print(f'accuracy before: {accuracy(model(x_train), y_train)}')\n",
        "\n",
        "    for i_epoch in range(num_epochs):\n",
        "        i_batch_start = 0\n",
        "        while i_batch_start + batch_size < x_train.shape[0]:\n",
        "            x_b = x_train[i_batch_start:i_batch_start + batch_size]\n",
        "            y_b = y_train[i_batch_start:i_batch_start + batch_size]\n",
        "            i_batch_start += batch_size\n",
        "\n",
        "            logit_b = model(x_b)    # an nn.Module is callable\n",
        "            loss = F.cross_entropy(logit_b, y_b)\n",
        "            loss.backward()\n",
        "\n",
        "            with torch.no_grad():\n",
        "                for p in model.parameters():\n",
        "                    p -= lr * p.grad\n",
        "                model.zero_grad()\n",
        "\n",
        "    print(f'accuracy after: {accuracy(model(x_train), y_train)}')\n",
        "\n",
        "train(MNISTLinear())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "czAS_2oo4qrP"
      },
      "source": [
        "`nn` contains many `nn.Module` subclasses that correspond to commonly used chunks of computation. One such \"layer\" is the linear model we've been using! We can replace our `MNISTLinear` class with an `nn.Linear`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ed2QkQ768CLI",
        "outputId": "ee83c159-44c4-4c4f-9f07-b2c49e3e630b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "accuracy before: 0.06741999834775925\n",
            "accuracy after: 0.9120399951934814\n"
          ]
        }
      ],
      "source": [
        "train(nn.Linear(784, 10))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B2qhStDy8lrP"
      },
      "source": [
        "Our optimization went differently in 2 epochs this time around. This is probably because the parameters in `nn.Linear` are [initialized slightly differently](https://pytorch.org/docs/stable/generated/torch.nn.Linear.html). It's always important to read documentation!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hkOLJWAH8zYD"
      },
      "source": [
        "### MNIST MLP\n",
        "What if we want to use a two-layer fully-connected network? `nn.Module`s can register `nn.Module` attributes!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "60V-yAKM5m00",
        "outputId": "5c11026a-c3db-4cca-8ef3-a95bd858a3f3"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "accuracy before: 0.12398000061511993\n",
            "accuracy after: 0.9607599973678589\n"
          ]
        }
      ],
      "source": [
        "class MNISTMLP(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.linear1 = nn.Linear(784, 64)\n",
        "        self.activation = nn.ReLU()\n",
        "        self.linear2 = nn.Linear(64, 10)\n",
        "\n",
        "    def forward(self, x_b):\n",
        "        return self.linear2(self.activation(self.linear1(x_b)))\n",
        "\n",
        "train(MNISTMLP())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eG2yPUjhYzzx"
      },
      "source": [
        "##Saving/Loading pre-trained weights\n",
        "Sometimes multiple people may have to use the same network. It can take a lot of time for everyone to train their network from scratch. It is environmentally also disastrous as the carbon footprint of training large networks is very huge.\n",
        "\n",
        "It would be better if we could train and save weights for our network, and reuse them later if we need to. PyTorch gives us ways to do this easily!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BjZ2iBpbbPDS",
        "outputId": "0d9f9587-6ef8-4063-e6fc-52fdf33ebc6c"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "accuracy before: 0.07614000141620636\n",
            "accuracy after: 0.953819990158081\n"
          ]
        }
      ],
      "source": [
        "network = MNISTMLP()\n",
        "train(network)\n",
        "\n",
        "torch.save(network.state_dict(), \"network.pt\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9Ew8Ken2b4He"
      },
      "source": [
        "The above code saves the network's weights to the given path, in this case \"network.pt\". Now we will show that loading the model's weight without retraining it from scratch retrieves prior performance."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4doKE58-cJV3",
        "outputId": "33298a9a-f71f-437f-9238-1fb4b933254e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "accuracy of initialized network: 0.1163799986243248\n",
            "accuracy after loading weights: 0.953819990158081\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "<ipython-input-22-2d02430dcb42>:4: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
            "  network.load_state_dict(torch.load(\"network.pt\"))\n"
          ]
        }
      ],
      "source": [
        "network = MNISTMLP()\n",
        "print(f'accuracy of initialized network: {accuracy(network(x_train), y_train)}')\n",
        "\n",
        "network.load_state_dict(torch.load(\"network.pt\"))\n",
        "print(f'accuracy after loading weights: {accuracy(network(x_train), y_train)}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FYXTLH477ooG"
      },
      "source": [
        "### Are You Registered?\n",
        "`nn.Module`s [have some specific rules](https://pytorch.org/docs/stable/generated/torch.nn.Module.html) regarding what attributes are properly registered (i.e. identified as part of the model). We've seen that `nn.Parameter`s and `nn.Module`s are, but most Python objects are not. For example, you might consider it cleaner to implement the `MNISTMLP` this way:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8sIQ_XOf-NM1",
        "outputId": "193b3ab6-9e85-4cca-bf94-4c188a85d638"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "accuracy before: 0.07054000347852707\n",
            "accuracy after: 0.07054000347852707\n",
            "[]\n"
          ]
        }
      ],
      "source": [
        "class MNISTMLP(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.layers = [\n",
        "            nn.Linear(784, 64),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(64, 10)\n",
        "        ]\n",
        "\n",
        "    def forward(self, x_b):\n",
        "        for layer in self.layers:\n",
        "            x_b = layer(x_b)\n",
        "        return x_b\n",
        "\n",
        "train(MNISTMLP())\n",
        "print(list(MNISTMLP().parameters()))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mn3Ogmn4-umv"
      },
      "source": [
        "Uh oh! Even though the `layers` list has `nn.Module`s, it itself wasn't recognized as a registerable attribute. To get around this, we have the `nn.ModuleList`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IA21B41t-7fL",
        "outputId": "69e60dfc-fce2-4169-a68d-4b51058d5e33"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "accuracy before: 0.10694000124931335\n",
            "accuracy after: 0.9515399932861328\n"
          ]
        }
      ],
      "source": [
        "class MNISTMLP(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.layers = nn.ModuleList([\n",
        "            nn.Linear(784, 64),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(64, 10)\n",
        "        ])\n",
        "\n",
        "    def forward(self, x_b):\n",
        "        for layer in self.layers:\n",
        "            x_b = layer(x_b)\n",
        "        return x_b\n",
        "\n",
        "train(MNISTMLP())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sN4KUQfu_QW7"
      },
      "source": [
        "### `nn.Sequential`\n",
        "It's pretty common to have a sequence of layers we want to iteratively apply onto an input. We can replace the `for` loop business in the forward pass by using an `nn.Sequential` container, a very commonly used abstraction:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rooDuNM2_C4S",
        "outputId": "9d447583-da76-4768-aec0-5e2a67707f62"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "accuracy before: 0.07383999973535538\n",
            "accuracy after: 0.9613400101661682\n"
          ]
        }
      ],
      "source": [
        "train(nn.Sequential(\n",
        "    nn.Linear(784, 64),\n",
        "    nn.ReLU(),\n",
        "    nn.Linear(64, 10)\n",
        "))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Zx745oji_vrm"
      },
      "source": [
        "### `torch.optim`\n",
        "`torch.optim` contains implementations of many common optimization algorithms. When initializing an optimizer instance, we need to pass in the model parameters. Now, instead of manually coding the update rule, we call the optimizer's `step` method. The optimizer can also handle resetting parameter gradients."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tupyWULJ7ZpU",
        "outputId": "dcc4a2e1-3d6e-449e-9ebf-68906fc70af9"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "accuracy before: 0.09752000123262405\n",
            "accuracy after: 0.9586799740791321\n"
          ]
        }
      ],
      "source": [
        "from torch import optim\n",
        "\n",
        "def train(model):\n",
        "    optimizer = optim.SGD(model.parameters(), lr=lr)\n",
        "\n",
        "    print(f'accuracy before: {accuracy(model(x_train), y_train)}')\n",
        "\n",
        "    for i_epoch in range(num_epochs):\n",
        "        i_batch_start = 0\n",
        "        while i_batch_start + batch_size < x_train.shape[0]:\n",
        "            x_b = x_train[i_batch_start:i_batch_start + batch_size]\n",
        "            y_b = y_train[i_batch_start:i_batch_start + batch_size]\n",
        "            i_batch_start += batch_size\n",
        "\n",
        "            logit_b = model(x_b)\n",
        "            loss = F.cross_entropy(logit_b, y_b)\n",
        "            loss.backward()\n",
        "\n",
        "            optimizer.step()\n",
        "            optimizer.zero_grad()\n",
        "\n",
        "    print(f'accuracy after: {accuracy(model(x_train), y_train)}')\n",
        "\n",
        "train(MNISTMLP())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h7SRGPQQ-VSB"
      },
      "source": [
        "### `torch.utils.data`\n",
        "Notice how we are manually slicing the `x_train` and `y_train` tensors to get batches. There's a more convenient way to do this:\n",
        "`torch.utils.data` features the `Dataset`, `Sampler`, and `DataLoader` abstractions, which are useful for transforming, sampling, and iterating over data. Let's implement a custom `Dataset` subclass for MNIST which we can then use to instantiate a `DataLoader`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wL0iPSSRKpER",
        "outputId": "8e95e611-b21f-4a85-8231-3dc0263907a4"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "accuracy before: 0.09415999799966812\n",
            "accuracy after: 0.9545800089836121\n"
          ]
        }
      ],
      "source": [
        "from torch.utils import data\n",
        "\n",
        "class MNISTDataset(data.Dataset):\n",
        "    def __init__(self, x, y):\n",
        "        self.x = x\n",
        "        self.y = y\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.y)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        return self.x[idx], self.y[idx]\n",
        "\n",
        "dataset_train = MNISTDataset(x_train, y_train)\n",
        "dataloader_train = data.DataLoader(dataset_train, batch_size=batch_size)\n",
        "\n",
        "def train(model):\n",
        "    optimizer = optim.SGD(model.parameters(), lr=lr)\n",
        "\n",
        "    print(f'accuracy before: {accuracy(model(x_train), y_train)}')\n",
        "\n",
        "    for i_epoch in range(num_epochs):\n",
        "        for x_b, y_b in dataloader_train:   # dataloaders are iterators\n",
        "            logit_b = model(x_b)\n",
        "            loss = F.cross_entropy(logit_b, y_b)\n",
        "            loss.backward()\n",
        "\n",
        "            optimizer.step()\n",
        "            optimizer.zero_grad()\n",
        "\n",
        "    print(f'accuracy after: {accuracy(model(x_train), y_train)}')\n",
        "\n",
        "train(MNISTMLP())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J3flJEs4LKZK"
      },
      "source": [
        "### Using a GPU\n",
        "Now, let's see how to use a GPU in PyTorch. Every tensor has a `device` attribute (default: `'cpu'`). GPU-accelerated computation requires the involved tensors to all be on device `'cuda'`. The `to` method is a handy way to move tensors around.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iWosupXrT_IV",
        "outputId": "2291cab1-6416-435b-fddd-3c902ef239b2"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "cpu\n",
            "cuda:0\n"
          ]
        }
      ],
      "source": [
        "x = torch.tensor(1.)\n",
        "print(x.device)\n",
        "x = x.to('cuda')\n",
        "print(x.device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E6ms19-qYfHB"
      },
      "source": [
        "The `nn.Module` also has a `to` method that transfers all parameters (and buffer items) to a specified device."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EeokNbD3YDXr",
        "outputId": "b6b0f0a8-41ff-4f25-9964-a89e12a1a2f3"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "accuracy before: 0.11267999559640884\n",
            "accuracy after: 0.9650399684906006\n"
          ]
        }
      ],
      "source": [
        "from torch.utils import data\n",
        "\n",
        "dataset_train = data.TensorDataset(x_train, y_train)\n",
        "dataloader_train = data.DataLoader(dataset_train, batch_size=batch_size)\n",
        "\n",
        "def train(model):\n",
        "    optimizer = optim.SGD(model.parameters(), lr=lr)\n",
        "\n",
        "    print(f'accuracy before: {accuracy(model(x_train.to(\"cuda\")), y_train.to(\"cuda\"))}')\n",
        "\n",
        "    for i_epoch in range(num_epochs):\n",
        "        for x_b, y_b in dataloader_train:\n",
        "            x_b, y_b = x_b.to('cuda'), y_b.to('cuda')\n",
        "            logit_b = model(x_b)\n",
        "            loss = F.cross_entropy(logit_b, y_b)\n",
        "            loss.backward()\n",
        "\n",
        "            optimizer.step()\n",
        "            optimizer.zero_grad()\n",
        "\n",
        "    print(f'accuracy after: {accuracy(model(x_train.to(\"cuda\")), y_train.to(\"cuda\"))}')\n",
        "\n",
        "train(MNISTMLP().to('cuda'))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eIKUTPPTVT2e"
      },
      "source": [
        "### Clean-up\n",
        "To wrap this section up, we'll refactor our code to facilitate assessing on both train and val data with data-loading, and using the GPU."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "PCtv2mcTAWSF",
        "outputId": "7e7f00bc-f0a0-4bc0-8756-2134c046a47c"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "before\n",
            "train loss and acc: (2.3162325, 0.07892)\n",
            "val loss and acc: (2.315925, 0.0786)\n",
            "after\n",
            "train loss and acc: (0.150307724609375, 0.95222)\n",
            "val loss and acc: (0.1640982666015625, 0.9514)\n"
          ]
        }
      ],
      "source": [
        "dataset_val = data.TensorDataset(x_val, y_val)\n",
        "dataloader_val = data.DataLoader(dataset_val, batch_size=batch_size)\n",
        "\n",
        "def loss_fn(model, x_b, y_b):\n",
        "    logit_b = model(x_b)\n",
        "    loss = F.cross_entropy(logit_b, y_b)\n",
        "    return loss, y_b.shape[0]\n",
        "\n",
        "def accuracy(model, x_b, y_b):\n",
        "    logit_b = model(x_b)\n",
        "    accuracy = (logit_b.argmax(dim=1) == y_b).float().mean()\n",
        "    return accuracy, y_b.shape[0]\n",
        "\n",
        "def assess(model, dataloader):\n",
        "    loss_total_r, correct_r, n_r = 0, 0, 0\n",
        "    with torch.no_grad():\n",
        "        for x_b, y_b in dataloader:\n",
        "            x_b, y_b = x_b.to('cuda'), y_b.to('cuda')\n",
        "            loss, n = loss_fn(model, x_b, y_b)\n",
        "            acc, n = accuracy(model, x_b, y_b)\n",
        "            loss_total_r += loss * n\n",
        "            correct_r += acc * n\n",
        "            n_r += n\n",
        "    loss_total_r, correct_r = loss_total_r.item(), correct_r.item()\n",
        "    return loss_total_r / n_r, correct_r / n_r\n",
        "\n",
        "def train(model, optimizer):\n",
        "    print('before')\n",
        "    print(f'train loss and acc: {assess(model, dataloader_train)}')\n",
        "    print(f'val loss and acc: {assess(model, dataloader_val)}')\n",
        "\n",
        "    for i_epoch in range(num_epochs):\n",
        "        for x_b, y_b in dataloader_train:\n",
        "            x_b, y_b = x_b.to('cuda'), y_b.to('cuda')\n",
        "            loss, _ = loss_fn(model, x_b, y_b)\n",
        "            loss.backward()\n",
        "\n",
        "            optimizer.step()\n",
        "            optimizer.zero_grad()\n",
        "\n",
        "    print('after')\n",
        "    print(f'train loss and acc: {assess(model, dataloader_train)}')\n",
        "    print(f'val loss and acc: {assess(model, dataloader_val)}')\n",
        "\n",
        "model = MNISTMLP().to('cuda')\n",
        "optimizer = optim.SGD(model.parameters(), lr=lr)\n",
        "train(model, optimizer)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BwUn0NLuVjeY"
      },
      "source": [
        "# More PyTorch Features: A Whirlwind Tour"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xkZ98YSob72X"
      },
      "source": [
        "### MNIST CNN\n",
        "2D convolutional layers expect an input of shape $(B, C, H, W)$ or `(batch, channels, height, width)`. For a batch of MNIST images, this is $(B, 1, 28, 28)$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tfXPPG3gb_cz",
        "outputId": "d77d9068-5e7a-4f4b-e585-26900ccf1a1e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "before\n",
            "train loss and acc: (2.3030059375, 0.0961)\n",
            "val loss and acc: (2.30335703125, 0.0883)\n",
            "after\n",
            "train loss and acc: (0.03881703369140625, 0.9874)\n",
            "val loss and acc: (0.056931982421875, 0.9844)\n"
          ]
        }
      ],
      "source": [
        "class MNISTCNN(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.layers = nn.Sequential(\n",
        "            nn.Conv2d(in_channels=1, out_channels=20, kernel_size=(5, 5), stride=1, padding=0),\n",
        "            nn.MaxPool2d(kernel_size=2, stride=2, padding=0),\n",
        "            nn.Conv2d(in_channels=20, out_channels=50, kernel_size=(5, 5), stride=1, padding=0),\n",
        "            nn.MaxPool2d(kernel_size=2, stride=2, padding=0),\n",
        "            nn.Flatten(),\n",
        "            nn.Linear(800, 500),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(500, 10)\n",
        "        )\n",
        "\n",
        "    def forward(self, x_b):\n",
        "        return self.layers(x_b.view((-1, 1, 28, 28)))\n",
        "\n",
        "model = MNISTCNN().to('cuda')\n",
        "optimizer = optim.Adam(model.parameters(), lr=1e-3)\n",
        "train(model, optimizer)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Transformer and TorchText: Sequence-to-Sequence Modeling \n",
        "This is a tutorial on how to train a sequence-to-sequence model that uses the nn.Transformer module.\n",
        "\n",
        "PyTorch 1.2 release includes a standard transformer module based on the paper Attention is All You Need. The transformer model has been proved to be superior in quality for many sequence-to-sequence problems while being more parallelizable. The nn.Transformer module relies entirely on an attention mechanism (another module recently implemented as nn.MultiheadAttention) to draw global dependencies between input and output. The nn.Transformer module is now highly modularized such that a single component (like nn.TransformerEncoder in this tutorial) can be easily adapted/composed"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Define the model\n",
        "\n",
        "\n",
        "In this tutorial, we train nn.TransformerEncoder model on a language modeling task. The language modeling task is to assign a probability for the likelihood of a given word (or a sequence of words) to follow a sequence of words. A sequence of tokens are passed to the embedding layer first, followed by a positional encoding layer to account for the order of the word (see the next paragraph for more details). The nn.TransformerEncoder consists of multiple layers of nn.TransformerEncoderLayer. Along with the input sequence, a square attention mask is required because the self-attention layers in nn.TransformerEncoder are only allowed to attend the earlier positions in the sequence. For the language modeling task, any tokens on the future positions should be masked. To have the actual words, the output of nn.TransformerEncoder model is sent to the final Linear layer, which is followed by a log-Softmax function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import math\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "class TransformerModel(nn.Module):\n",
        "\n",
        "    def __init__(self, ntoken, ninp, nhead, nhid, nlayers, dropout=0.5):\n",
        "        super(TransformerModel, self).__init__()\n",
        "        from torch.nn import TransformerEncoder, TransformerEncoderLayer\n",
        "        self.model_type = 'Transformer'\n",
        "        self.src_mask = None\n",
        "        self.pos_encoder = PositionalEncoding(ninp, dropout)\n",
        "        encoder_layers = TransformerEncoderLayer(ninp, nhead, nhid, dropout)\n",
        "        self.transformer_encoder = TransformerEncoder(encoder_layers, nlayers)\n",
        "        self.encoder = nn.Embedding(ntoken, ninp)\n",
        "        self.ninp = ninp\n",
        "        self.decoder = nn.Linear(ninp, ntoken)\n",
        "\n",
        "        self.init_weights()\n",
        "\n",
        "    def _generate_square_subsequent_mask(self, sz):\n",
        "        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)\n",
        "        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))\n",
        "        return mask\n",
        "\n",
        "    def init_weights(self):\n",
        "        initrange = 0.1\n",
        "        self.encoder.weight.data.uniform_(-initrange, initrange)\n",
        "        self.decoder.bias.data.zero_()\n",
        "        self.decoder.weight.data.uniform_(-initrange, initrange)\n",
        "\n",
        "    def forward(self, src):\n",
        "        if self.src_mask is None or self.src_mask.size(0) != len(src):\n",
        "            device = src.device\n",
        "            mask = self._generate_square_subsequent_mask(len(src)).to(device)\n",
        "            self.src_mask = mask\n",
        "\n",
        "        src = self.encoder(src) * math.sqrt(self.ninp)\n",
        "        src = self.pos_encoder(src)\n",
        "        output = self.transformer_encoder(src, self.src_mask)\n",
        "        output = self.decoder(output)\n",
        "        return output"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "PositionalEncoding module injects some information about the relative or absolute position of the tokens in the sequence. The positional encodings have the same dimension as the embeddings so that the two can be summed. Here, we use sine and cosine functions of different frequencies."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "class PositionalEncoding(nn.Module):\n",
        "\n",
        "    def __init__(self, d_model, dropout=0.1, max_len=5000):\n",
        "        super(PositionalEncoding, self).__init__()\n",
        "        self.dropout = nn.Dropout(p=dropout)\n",
        "\n",
        "        pe = torch.zeros(max_len, d_model)\n",
        "        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)\n",
        "        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))\n",
        "        pe[:, 0::2] = torch.sin(position * div_term)\n",
        "        pe[:, 1::2] = torch.cos(position * div_term)\n",
        "        pe = pe.unsqueeze(0).transpose(0, 1)\n",
        "        self.register_buffer('pe', pe)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x + self.pe[:x.size(0), :]\n",
        "        return self.dropout(x)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Load and batch data\n",
        "The training process uses Wikitext-2 dataset from torchtext. The vocab object is built based on the train dataset and is used to numericalize tokens into tensors. Starting from sequential data, the batchify() function arranges the dataset into columns, trimming off any tokens remaining after the data has been divided into batches of size batch_size. For instance, with the alphabet as the sequence (total length of 26) and a batch size of 4, we would divide the alphabet into 4 sequences of length 6:\n",
        "\n",
        "These columns are treated as independent by the model, which means that the dependence of G and F can not be learned, but allows more efficient batch processing."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import torchtext\n",
        "from torchtext.data.utils import get_tokenizer\n",
        "TEXT = torchtext.data.Field(tokenize=get_tokenizer(\"basic_english\"),\n",
        "                            init_token='<sos>',\n",
        "                            eos_token='<eos>',\n",
        "                            lower=True)\n",
        "train_txt, val_txt, test_txt = torchtext.datasets.WikiText2.splits(TEXT)\n",
        "TEXT.build_vocab(train_txt)\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "def batchify(data, bsz):\n",
        "    data = TEXT.numericalize([data.examples[0].text])\n",
        "    # Divide the dataset into bsz parts.\n",
        "    nbatch = data.size(0) // bsz\n",
        "    # Trim off any extra elements that wouldn't cleanly fit (remainders).\n",
        "    data = data.narrow(0, 0, nbatch * bsz)\n",
        "    # Evenly divide the data across the bsz batches.\n",
        "    data = data.view(bsz, -1).t().contiguous()\n",
        "    return data.to(device)\n",
        "\n",
        "batch_size = 20\n",
        "eval_batch_size = 10\n",
        "train_data = batchify(train_txt, batch_size)\n",
        "val_data = batchify(val_txt, eval_batch_size)\n",
        "test_data = batchify(test_txt, eval_batch_size)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Out: "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "downloading wikitext-2-v1.zip\n",
        "extracting"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Functions to generate input and target sequence\n",
        "get_batch() function generates the input and target sequence for the transformer model. It subdivides the source data into chunks of length bptt. For the language modeling task, the model needs the following words as Target. For example, with a bptt value of 2, we’d get the following two Variables for i = 0:\n",
        "\n",
        "It should be noted that the chunks are along dimension 0, consistent with the S dimension in the Transformer model. The batch dimension N is along dimension 1.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "bptt = 35\n",
        "def get_batch(source, i):\n",
        "    seq_len = min(bptt, len(source) - 1 - i)\n",
        "    data = source[i:i+seq_len]\n",
        "    target = source[i+1:i+1+seq_len].view(-1)\n",
        "    return data, target"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Initiate an instance\n",
        "The model is set up with the hyperparameter below. The vocab size is equal to the length of the vocab object."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "ntokens = len(TEXT.vocab.stoi) # the size of vocabulary\n",
        "emsize = 200 # embedding dimension\n",
        "nhid = 200 # the dimension of the feedforward network model in nn.TransformerEncoder\n",
        "nlayers = 2 # the number of nn.TransformerEncoderLayer in nn.TransformerEncoder\n",
        "nhead = 2 # the number of heads in the multiheadattention models\n",
        "dropout = 0.2 # the dropout value\n",
        "model = TransformerModel(ntokens, emsize, nhead, nhid, nlayers, dropout).to(device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Run the model\n",
        "\n",
        "CrossEntropyLoss is applied to track the loss and SGD implements stochastic gradient descent method as the optimizer. The initial learning rate is set to 5.0. StepLR is applied to adjust the learn rate through epochs. During the training, we use nn.utils.clip_grad_norm_ function to scale all the gradient together to prevent exploding."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "criterion = nn.CrossEntropyLoss()\n",
        "lr = 5.0 # learning rate\n",
        "optimizer = torch.optim.SGD(model.parameters(), lr=lr)\n",
        "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.95)\n",
        "\n",
        "import time\n",
        "def train():\n",
        "    model.train() # Turn on the train mode\n",
        "    total_loss = 0.\n",
        "    start_time = time.time()\n",
        "    ntokens = len(TEXT.vocab.stoi)\n",
        "    for batch, i in enumerate(range(0, train_data.size(0) - 1, bptt)):\n",
        "        data, targets = get_batch(train_data, i)\n",
        "        optimizer.zero_grad()\n",
        "        output = model(data)\n",
        "        loss = criterion(output.view(-1, ntokens), targets)\n",
        "        loss.backward()\n",
        "        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)\n",
        "        optimizer.step()\n",
        "\n",
        "        total_loss += loss.item()\n",
        "        log_interval = 200\n",
        "        if batch % log_interval == 0 and batch > 0:\n",
        "            cur_loss = total_loss / log_interval\n",
        "            elapsed = time.time() - start_time\n",
        "            print('| epoch {:3d} | {:5d}/{:5d} batches | '\n",
        "                  'lr {:02.2f} | ms/batch {:5.2f} | '\n",
        "                  'loss {:5.2f} | ppl {:8.2f}'.format(\n",
        "                    epoch, batch, len(train_data) // bptt, scheduler.get_lr()[0],\n",
        "                    elapsed * 1000 / log_interval,\n",
        "                    cur_loss, math.exp(cur_loss)))\n",
        "            total_loss = 0\n",
        "            start_time = time.time()\n",
        "\n",
        "def evaluate(eval_model, data_source):\n",
        "    eval_model.eval() # Turn on the evaluation mode\n",
        "    total_loss = 0.\n",
        "    ntokens = len(TEXT.vocab.stoi)\n",
        "    with torch.no_grad():\n",
        "        for i in range(0, data_source.size(0) - 1, bptt):\n",
        "            data, targets = get_batch(data_source, i)\n",
        "            output = eval_model(data)\n",
        "            output_flat = output.view(-1, ntokens)\n",
        "            total_loss += len(data) * criterion(output_flat, targets).item()\n",
        "    return total_loss / (len(data_source) - 1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Loop over epochs. Save the model if the validation loss is the best we’ve seen so far. Adjust the learning rate after each epoch."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "best_val_loss = float(\"inf\")\n",
        "epochs = 3 # The number of epochs\n",
        "best_model = None\n",
        "\n",
        "for epoch in range(1, epochs + 1):\n",
        "    epoch_start_time = time.time()\n",
        "    train()\n",
        "    val_loss = evaluate(model, val_data)\n",
        "    print('-' * 89)\n",
        "    print('| end of epoch {:3d} | time: {:5.2f}s | valid loss {:5.2f} | '\n",
        "          'valid ppl {:8.2f}'.format(epoch, (time.time() - epoch_start_time),\n",
        "                                     val_loss, math.exp(val_loss)))\n",
        "    print('-' * 89)\n",
        "\n",
        "    if val_loss < best_val_loss:\n",
        "        best_val_loss = val_loss\n",
        "        best_model = model\n",
        "\n",
        "    scheduler.step()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Out:\n",
        "\n",
        "| epoch   1 |   200/ 2981 batches | lr 5.00 | ms/batch 29.39 | loss  8.05 | ppl  3123.83\n",
        "| epoch   1 |   400/ 2981 batches | lr 5.00 | ms/batch 28.52 | loss  6.74 | ppl   844.07\n",
        "| epoch   1 |   600/ 2981 batches | lr 5.00 | ms/batch 28.50 | loss  6.35 | ppl   571.65\n",
        "| epoch   1 |   800/ 2981 batches | lr 5.00 | ms/batch 28.60 | loss  6.22 | ppl   505.09\n",
        "| epoch   1 |  1000/ 2981 batches | lr 5.00 | ms/batch 28.60 | loss  6.11 | ppl   451.99\n",
        "| epoch   1 |  1200/ 2981 batches | lr 5.00 | ms/batch 28.57 | loss  6.08 | ppl   439.11\n",
        "| epoch   1 |  1400/ 2981 batches | lr 5.00 | ms/batch 28.60 | loss  6.04 | ppl   419.65\n",
        "| epoch   1 |  1600/ 2981 batches | lr 5.00 | ms/batch 28.53 | loss  6.05 | ppl   426.08\n",
        "| epoch   1 |  1800/ 2981 batches | lr 5.00 | ms/batch 28.60 | loss  5.96 | ppl   387.40\n",
        "| epoch   1 |  2000/ 2981 batches | lr 5.00 | ms/batch 28.61 | loss  5.96 | ppl   387.10\n",
        "| epoch   1 |  2200/ 2981 batches | lr 5.00 | ms/batch 28.61 | loss  5.85 | ppl   346.65\n",
        "| epoch   1 |  2400/ 2981 batches | lr 5.00 | ms/batch 28.59 | loss  5.90 | ppl   363.66\n",
        "| epoch   1 |  2600/ 2981 batches | lr 5.00 | ms/batch 28.58 | loss  5.89 | ppl   362.43\n",
        "| epoch   1 |  2800/ 2981 batches | lr 5.00 | ms/batch 28.63 | loss  5.79 | ppl   327.16\n",
        "-----------------------------------------------------------------------------------------\n",
        "| end of epoch   1 | time: 88.73s | valid loss  5.69 | valid ppl   297.35\n",
        "-----------------------------------------------------------------------------------------\n",
        "| epoch   2 |   200/ 2981 batches | lr 4.51 | ms/batch 28.74 | loss  5.79 | ppl   328.47\n",
        "| epoch   2 |   400/ 2981 batches | lr 4.51 | ms/batch 28.63 | loss  5.77 | ppl   320.22\n",
        "| epoch   2 |   600/ 2981 batches | lr 4.51 | ms/batch 28.60 | loss  5.60 | ppl   270.08\n",
        "| epoch   2 |   800/ 2981 batches | lr 4.51 | ms/batch 28.58 | loss  5.64 | ppl   280.65\n",
        "| epoch   2 |  1000/ 2981 batches | lr 4.51 | ms/batch 28.61 | loss  5.59 | ppl   267.77\n",
        "| epoch   2 |  1200/ 2981 batches | lr 4.51 | ms/batch 28.51 | loss  5.61 | ppl   272.92\n",
        "| epoch   2 |  1400/ 2981 batches | lr 4.51 | ms/batch 28.58 | loss  5.62 | ppl   276.17\n",
        "| epoch   2 |  1600/ 2981 batches | lr 4.51 | ms/batch 28.52 | loss  5.66 | ppl   286.68\n",
        "| epoch   2 |  1800/ 2981 batches | lr 4.51 | ms/batch 28.62 | loss  5.59 | ppl   267.02\n",
        "| epoch   2 |  2000/ 2981 batches | lr 4.51 | ms/batch 28.65 | loss  5.61 | ppl   272.22\n",
        "| epoch   2 |  2200/ 2981 batches | lr 4.51 | ms/batch 28.60 | loss  5.51 | ppl   247.40\n",
        "| epoch   2 |  2400/ 2981 batches | lr 4.51 | ms/batch 28.55 | loss  5.57 | ppl   263.63\n",
        "| epoch   2 |  2600/ 2981 batches | lr 4.51 | ms/batch 28.55 | loss  5.59 | ppl   268.08\n",
        "| epoch   2 |  2800/ 2981 batches | lr 4.51 | ms/batch 28.63 | loss  5.51 | ppl   247.61\n",
        "-----------------------------------------------------------------------------------------\n",
        "| end of epoch   2 | time: 88.63s | valid loss  5.56 | valid ppl   259.39\n",
        "-----------------------------------------------------------------------------------------\n",
        "| epoch   3 |   200/ 2981 batches | lr 4.29 | ms/batch 28.74 | loss  5.55 | ppl   257.69\n",
        "| epoch   3 |   400/ 2981 batches | lr 4.29 | ms/batch 28.56 | loss  5.55 | ppl   256.46\n",
        "| epoch   3 |   600/ 2981 batches | lr 4.29 | ms/batch 28.60 | loss  5.35 | ppl   211.22\n",
        "| epoch   3 |   800/ 2981 batches | lr 4.29 | ms/batch 28.58 | loss  5.42 | ppl   226.08\n",
        "| epoch   3 |  1000/ 2981 batches | lr 4.29 | ms/batch 28.62 | loss  5.38 | ppl   216.52\n",
        "| epoch   3 |  1200/ 2981 batches | lr 4.29 | ms/batch 28.59 | loss  5.41 | ppl   223.14\n",
        "| epoch   3 |  1400/ 2981 batches | lr 4.29 | ms/batch 28.62 | loss  5.43 | ppl   228.66\n",
        "| epoch   3 |  1600/ 2981 batches | lr 4.29 | ms/batch 28.64 | loss  5.47 | ppl   238.30\n",
        "| epoch   3 |  1800/ 2981 batches | lr 4.29 | ms/batch 28.48 | loss  5.40 | ppl   221.59\n",
        "| epoch   3 |  2000/ 2981 batches | lr 4.29 | ms/batch 28.62 | loss  5.44 | ppl   230.29\n",
        "| epoch   3 |  2200/ 2981 batches | lr 4.29 | ms/batch 28.60 | loss  5.32 | ppl   205.12\n",
        "| epoch   3 |  2400/ 2981 batches | lr 4.29 | ms/batch 28.58 | loss  5.39 | ppl   219.94\n",
        "| epoch   3 |  2600/ 2981 batches | lr 4.29 | ms/batch 28.66 | loss  5.41 | ppl   224.03\n",
        "| epoch   3 |  2800/ 2981 batches | lr 4.29 | ms/batch 28.59 | loss  5.34 | ppl   208.79\n",
        "-----------------------------------------------------------------------------------------\n",
        "| end of epoch   3 | time: 88.64s | valid loss  5.55 | valid ppl   257.39\n",
        "-----------------------------------------------------------------------------------------"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Evaluate the model with the test dataset\n",
        "Apply the best model to check the result with the test dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "test_loss = evaluate(best_model, test_data)\n",
        "print('=' * 89)\n",
        "print('| End of training | test loss {:5.2f} | test ppl {:8.2f}'.format(\n",
        "    test_loss, math.exp(test_loss)))\n",
        "print('=' * 89)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Out:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "=========================================================================================\n",
        "| End of training | test loss  5.46 | test ppl   235.14\n",
        "========================================================================================="
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XLV7dSnk_xC4"
      },
      "source": [
        "### MNIST LSTM\n",
        "When constructed using the `batch_first=True` argument, `nn.LSTM`s expect inputs of shape $(B, L, D)$ or `(batch, sequence length, element size)`. We'll interpret the rows of the MNIST image as sequence elements, and compute the logits from the last LSTM output (i.e. wait for the LSTM to see the entire image)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CSA3chD5Aiu_",
        "outputId": "fb1dc73e-e1be-4ee7-9bbe-dbe3b1fc25b1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "before\n",
            "train loss and acc: (2.30728546875, 0.09976)\n",
            "val loss and acc: (2.30689140625, 0.0961)\n",
            "after\n",
            "train loss and acc: (0.17556994140625, 0.95036)\n",
            "val loss and acc: (0.164346044921875, 0.9524)\n"
          ]
        }
      ],
      "source": [
        "class MNISTLSTM(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.LSTM = nn.LSTM(input_size=28, hidden_size=64, num_layers=1, batch_first=True)\n",
        "        self.linear = nn.Linear(64, 10)\n",
        "\n",
        "    def forward(self, x_b):\n",
        "        out, _ = self.LSTM(x_b.view((-1, 28, 28)))  # out contains outputs at each iteration over the sequence\n",
        "        return self.linear(out[:, -1, :])   # only the out at the last iteration has seen the entire image\n",
        "\n",
        "model = MNISTLSTM().to('cuda')\n",
        "optimizer = optim.Adam(model.parameters(), lr=1e-3)\n",
        "train(model, optimizer)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j7gNjS2HEeEo"
      },
      "source": [
        "### `nn.Embedding`\n",
        "Oftentimes we will have categorical data with too many categories to comfortably convert into one-hot encodings. Examples of this include tokens in NLP and $(x,y)$ positions on a grid. One common way to process this data is to use an `nn.Embedding`, or a table of dense vectors that can be indexed by the categorical data. The dense vectors can be optimized. `nn.Embedding` supports advanced indexing."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_45NPu8MHZMu",
        "outputId": "88a72442-b288-430b-b3ae-0fc5bf4f5484"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "torch.Size([64])\n",
            "True\n",
            "torch.Size([2, 4])\n",
            "torch.Size([2, 4, 64])\n"
          ]
        }
      ],
      "source": [
        "embedding = nn.Embedding(num_embeddings=1000, embedding_dim=64)\n",
        "print(embedding(torch.tensor(42)).shape)\n",
        "print(embedding(torch.tensor(42)).requires_grad)\n",
        "indices = torch.tensor([\n",
        "    [0, 1, 42, 999],\n",
        "    [1, 1, 2, 3]\n",
        "])\n",
        "print(indices.shape)\n",
        "print(embedding(indices).shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G4hyCvaoIq-c"
      },
      "source": [
        "If you're interested in NLP, a nice use of `nn.Embedding` and recurrent layers (`nn.GRU`) for natural language translation can be found in [this tutorial](https://pytorch.org/tutorials/intermediate/seq2seq_translation_tutorial.html)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pun7z9zGqbvH"
      },
      "source": [
        "### `torch.distributions`\n",
        "Aside from conveniently providing commonly used methods (e.g. density/mass evaluation), one of the main features of this module is the use of the reparameterization trick to facilitate backpropagation through random samples to the underlying distribution parameters (and beyond)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BHyPfnxdJvc4",
        "outputId": "ea71fe39-676b-4fa5-ea55-f008d9d2b026"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor([ 0.3411,  0.0985, -0.5268,  0.2175, -2.2333], grad_fn=<AddBackward0>)\n",
            "tensor(-2.1754)\n"
          ]
        }
      ],
      "source": [
        "from torch import distributions\n",
        "loc = torch.tensor(0., requires_grad=True)\n",
        "scale = torch.tensor(1., requires_grad=True)\n",
        "p = distributions.Normal(loc, scale)\n",
        "x = p.rsample(torch.Size([5]))\n",
        "print(x)    # grad_fn comes from reparameterization trick\n",
        "y = -torch.mean(x**2)\n",
        "y.backward()\n",
        "print(scale.grad)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vRZRNM6MLZMU"
      },
      "source": [
        "If we don't use `rsample`, we can't take gradients through the sampling step."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 38,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sWDu3fi6LN0y",
        "outputId": "fb515cde-3a40-4d33-93e9-a84651d29c15"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor([-0.3452, -1.1077, -0.3113, -2.0162,  0.3057], grad_fn=<AddBackward0>)\n"
          ]
        }
      ],
      "source": [
        "x = p.rsample(torch.Size([5]))\n",
        "print(x)\n",
        "y = -torch.sum(x**2)\n",
        "y.backward()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GrIbkiJhNyRF"
      },
      "source": [
        "# Debugging, Documentation, and Getting Help\n",
        "Your assignments won't be in Jupyter notebooks, but in Python scripts. This means you won't have the interactivity of cell-based execution, but to the rescue comes `pdb`, the Python Debugger. Use it to\n",
        "\n",
        "*   set breakpoints in the code to interactively inspect program elements: `import pdb; pdb.set_trace()`\n",
        "*   automatically start a debugging session when an exception is thrown: `python -m pdb -c continue main.py`\n",
        "\n",
        "Both will be super, super useful.\n",
        "\n",
        "When something about PyTorch is confusing you, the first place to look is the [PyTorch documentation](https://https://pytorch.org/docs/stable/index.html). (The author of this tutorial consulted the documentation no fewer than 20 times when making it.) This is often the last place you need to look, but sometimes the [PyTorch forums](https://https://discuss.pytorch.org/) can be helpful, too. (Use the search tool!) And of course, you're always welcome to post on Ed or seek out one of your friendly neighborhood CAs.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c5FIR5icqCax"
      },
      "source": [
        "\n",
        "\n",
        "---\n",
        "\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "include_colab_link": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}