Notebooks/Chap08/8_1_MNIST_1D_Performance.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/udlbook/udlbook/blob/main/Notebooks/Chap08/8_1_MNIST_1D_Performance.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L6chybAVFJW2"
      },
      "source": [
        "# **Notebook 8.1: MNIST_1D_Performance**\n",
        "\n",
        "This notebook runs a simple neural network on the MNIST1D dataset as in figure 8.2a. It uses code from https://github.com/greydanus/mnist1d to generate the data.\n",
        "\n",
        "Work through the cells below, running each cell in turn. In various places you will see the words \"TO DO\". Follow the instructions at these places and make predictions about what is going to happen or write code to complete the functions.\n",
        "\n",
        "Contact me at udlbookmail@gmail.com if you find any mistakes or have any suggestions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ifVjS4cTOqKz"
      },
      "outputs": [],
      "source": [
        "# Run this if you're in a Colab to install MNIST 1D repository\n",
        "%pip install git+https://github.com/greydanus/mnist1d"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qyE7G1StPIqO"
      },
      "outputs": [],
      "source": [
        "import torch, torch.nn as nn\n",
        "from torch.utils.data import TensorDataset, DataLoader\n",
        "from torch.optim.lr_scheduler import StepLR\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import mnist1d"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F7LNq72SP6jO"
      },
      "source": [
        "Let's generate a training and test dataset using the MNIST1D code.  The dataset gets saved as a .pkl file so it doesn't have to be regenerated each time."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YLxf7dJfPaqw"
      },
      "outputs": [],
      "source": [
        "args = mnist1d.data.get_dataset_args()\n",
        "data = mnist1d.data.get_dataset(args, path='./mnist1d_data.pkl', download=False, regenerate=False)\n",
        "\n",
        "# The training and test input and outputs are in\n",
        "# data['x'], data['y'], data['x_test'], and data['y_test']\n",
        "print(\"Examples in training set: {}\".format(len(data['y'])))\n",
        "print(\"Examples in test set: {}\".format(len(data['y_test'])))\n",
        "print(\"Length of each example: {}\".format(data['x'].shape[-1]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FxaB5vc0uevl"
      },
      "outputs": [],
      "source": [
        "D_i = 40    # Input dimensions\n",
        "D_k = 100   # Hidden dimensions\n",
        "D_o = 10    # Output dimensions\n",
        "# TO DO:\n",
        "# Define a model with two hidden layers of size 100\n",
        "# And ReLU activations between them\n",
        "# Replace this line (see Figure 7.8 of book for help):\n",
        "model = torch.nn.Sequential(torch.nn.Linear(D_i, D_o));\n",
        "\n",
        "\n",
        "def weights_init(layer_in):\n",
        "  # TO DO:\n",
        "  # Initialize the parameters with He initialization\n",
        "  # Replace this line (see figure 7.8 of book for help)\n",
        "  print(\"Initializing layer\")\n",
        "\n",
        "\n",
        "# Call the function you just defined\n",
        "model.apply(weights_init)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_rX6N3VyyQTY"
      },
      "outputs": [],
      "source": [
        "# choose cross entropy loss function (equation 5.24)\n",
        "loss_function = torch.nn.CrossEntropyLoss()\n",
        "# construct SGD optimizer and initialize learning rate and momentum\n",
        "optimizer = torch.optim.SGD(model.parameters(), lr = 0.05, momentum=0.9)\n",
        "# object that decreases learning rate by half every 10 epochs\n",
        "scheduler = StepLR(optimizer, step_size=10, gamma=0.5)\n",
        "x_train = torch.tensor(data['x'].astype('float32'))\n",
        "y_train = torch.tensor(data['y'].transpose().astype('int64'))\n",
        "x_test= torch.tensor(data['x_test'].astype('float32'))\n",
        "y_test = torch.tensor(data['y_test'].astype('int64'))\n",
        "\n",
        "# load the data into a class that creates the batches\n",
        "data_loader = DataLoader(TensorDataset(x_train,y_train), batch_size=100, shuffle=True, worker_init_fn=np.random.seed(1))\n",
        "\n",
        "# Initialize model weights\n",
        "model.apply(weights_init)\n",
        "\n",
        "# loop over the dataset n_epoch times\n",
        "n_epoch = 50\n",
        "# store the loss and the % correct at each epoch\n",
        "losses_train = np.zeros((n_epoch))\n",
        "errors_train = np.zeros((n_epoch))\n",
        "losses_test = np.zeros((n_epoch))\n",
        "errors_test = np.zeros((n_epoch))\n",
        "\n",
        "for epoch in range(n_epoch):\n",
        "  # loop over batches\n",
        "  for i, batch in enumerate(data_loader):\n",
        "    # retrieve inputs and labels for this batch\n",
        "    x_batch, y_batch = batch\n",
        "    # zero the parameter gradients\n",
        "    optimizer.zero_grad()\n",
        "    # forward pass -- calculate model output\n",
        "    pred = model(x_batch)\n",
        "    # compute the loss\n",
        "    loss = loss_function(pred, y_batch)\n",
        "    # backward pass\n",
        "    loss.backward()\n",
        "    # SGD update\n",
        "    optimizer.step()\n",
        "\n",
        "  # Run whole dataset to get statistics -- normally wouldn't do this\n",
        "  pred_train = model(x_train)\n",
        "  pred_test = model(x_test)\n",
        "  _, predicted_train_class = torch.max(pred_train.data, 1)\n",
        "  _, predicted_test_class = torch.max(pred_test.data, 1)\n",
        "  errors_train[epoch] = 100 - 100 * (predicted_train_class == y_train).float().sum() / len(y_train)\n",
        "  errors_test[epoch]= 100 - 100 * (predicted_test_class == y_test).float().sum() / len(y_test)\n",
        "  losses_train[epoch] = loss_function(pred_train, y_train).item()\n",
        "  losses_test[epoch]= loss_function(pred_test, y_test).item()\n",
        "  print(f'Epoch {epoch:5d}, train loss {losses_train[epoch]:.6f}, train error {errors_train[epoch]:3.2f},  test loss {losses_test[epoch]:.6f}, test error {errors_test[epoch]:3.2f}')\n",
        "\n",
        "  # tell scheduler to consider updating learning rate\n",
        "  scheduler.step()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yI-l6kA_EH9G"
      },
      "outputs": [],
      "source": [
        "# Plot the results\n",
        "fig, ax = plt.subplots()\n",
        "ax.plot(errors_train,'r-',label='train')\n",
        "ax.plot(errors_test,'b-',label='test')\n",
        "ax.set_ylim(0,100); ax.set_xlim(0,n_epoch)\n",
        "ax.set_xlabel('Epoch'); ax.set_ylabel('Error')\n",
        "ax.set_title('TrainError %3.2f, Test Error %3.2f'%(errors_train[-1],errors_test[-1]))\n",
        "ax.legend()\n",
        "plt.show()\n",
        "\n",
        "# Plot the results\n",
        "fig, ax = plt.subplots()\n",
        "ax.plot(losses_train,'r-',label='train')\n",
        "ax.plot(losses_test,'b-',label='test')\n",
        "ax.set_xlim(0,n_epoch)\n",
        "ax.set_xlabel('Epoch'); ax.set_ylabel('Loss')\n",
        "ax.set_title('Train loss %3.2f, Test loss %3.2f'%(losses_train[-1],losses_test[-1]))\n",
        "ax.legend()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q-yT6re6GZS4"
      },
      "source": [
        "**TO DO**\n",
        "\n",
        "Play with the model -- try changing the number of layers, hidden units, learning rate, batch size, momentum or anything else you like.  See if you can improve the test results.\n",
        "\n",
        "Is it a good idea to optimize the hyperparameters in this way?  Will the final result be a good estimate of the true test performance?"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "authorship_tag": "ABX9TyOuKMUcKfOIhIL2qTX9jJCy",
      "gpuType": "T4",
      "include_colab_link": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}