{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "mnist_with_pytorch.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z5achl3MnYXU"
      },
      "source": [
        "# MNIST With PyTorch Training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ncnRNj6wndZj"
      },
      "source": [
        "## Import Libraries"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nOB2NoPxlSep"
      },
      "source": [
        "import numpy as np\n",
        "import torch\n",
        "import torchvision\n",
        "import matplotlib.pyplot as plt\n",
        "from torchvision import datasets, transforms\n",
        "from torch import nn\n",
        "from torch import optim\n",
        "from time import time\n",
        "import os\n",
        "from google.colab import drive"
      ],
      "execution_count": 46,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nbt5QydhnfqH"
      },
      "source": [
        "## Pre-Process Data\n",
        "\n",
        "Here we download the data using PyTorch data utils and transform the data by using a normalization function. PyTorch provides a data loader abstraction called a `DataLoader` where we can set the batch size, data shuffle per batch loading. Each data loader expecte a Pytorch Dataset. The DataSet abstraction and DataLoader usage can be found [here](https://pytorch.org/tutorials/recipes/recipes/loading_data_recipe.html) "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yWKZO4sHlhSQ"
      },
      "source": [
        "# Data transformation function \n",
        "transform = transforms.Compose([transforms.ToTensor(),\n",
        "                              transforms.Normalize((0.5,), (0.5,)),\n",
        "                              ])\n",
        "\n",
        "# DataSet\n",
        "train_data_set = datasets.MNIST('drive/My Drive/mnist/data/', download=True, train=True, transform=transform)\n",
        "validation_data_set = datasets.MNIST('drive/My Drive/mnist/data/', download=True, train=False, transform=transform)\n",
        "\n",
        "\n",
        "train_loader = torch.utils.data.DataLoader(train_data_set, batch_size=32, shuffle=True)\n",
        "validation_loader = torch.utils.data.DataLoader(validation_data_set, batch_size=32, shuffle=True)"
      ],
      "execution_count": 47,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NXrmy7jYniQn"
      },
      "source": [
        "## Define Network\n",
        "\n",
        "Here we select the matching input size compared to the network definition. \n",
        "Here data reshaping or layer reshaping must be done to match input data shape with the network input shape. Also we define a set of hidden unit sizes along with the output layers size. The `output_size` must match with the number of labels associated with the classification problem. The hidden units can be chosesn depending on the problem. `nn.Sequential` is one way to create the network. Here we stack a set of linear layers along with a softmax layer for the classification as the output layer. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "k65r0ZeHlzgY",
        "outputId": "f78eba4c-144a-436b-888f-464c7cd95488"
      },
      "source": [
        "input_size = 784\n",
        "hidden_sizes = [128, 128, 64, 64]\n",
        "output_size = 10\n",
        "\n",
        "model = nn.Sequential(nn.Linear(input_size, hidden_sizes[0]),\n",
        "                      nn.ReLU(),\n",
        "                      nn.Linear(hidden_sizes[0], hidden_sizes[1]),\n",
        "                      nn.ReLU(),\n",
        "                      nn.Linear(hidden_sizes[1], hidden_sizes[2]),\n",
        "                      nn.ReLU(),\n",
        "                      nn.Linear(hidden_sizes[2], hidden_sizes[3]),\n",
        "                      nn.ReLU(),\n",
        "                      nn.Linear(hidden_sizes[3], output_size),\n",
        "                      nn.LogSoftmax(dim=1))\n",
        "\n",
        "                      \n",
        "print(model)"
      ],
      "execution_count": 48,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Sequential(\n",
            "  (0): Linear(in_features=784, out_features=128, bias=True)\n",
            "  (1): ReLU()\n",
            "  (2): Linear(in_features=128, out_features=128, bias=True)\n",
            "  (3): ReLU()\n",
            "  (4): Linear(in_features=128, out_features=64, bias=True)\n",
            "  (5): ReLU()\n",
            "  (6): Linear(in_features=64, out_features=64, bias=True)\n",
            "  (7): ReLU()\n",
            "  (8): Linear(in_features=64, out_features=10, bias=True)\n",
            "  (9): LogSoftmax(dim=1)\n",
            ")\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BzRdJkNJn3iB"
      },
      "source": [
        "## Define Loss Function and Optimizer\n",
        "\n",
        "Read more about [Loss Functions](https://pytorch.org/docs/stable/nn.html#loss-functions) and [Optimizers](https://pytorch.org/docs/stable/optim.html) supported by PyTorch.\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8BgdSHzyl1cM"
      },
      "source": [
        "criterion = nn.NLLLoss()\n",
        "optimizer = optim.SGD(model.parameters(), lr=0.003, momentum=0.9)"
      ],
      "execution_count": 49,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4BjYhjdFn9uG"
      },
      "source": [
        "# Model Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PgZDncI-mEoM",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "3205000a-7dd0-474c-bfea-6685b1b33859"
      },
      "source": [
        "epochs = 5\n",
        "\n",
        "for epoch in range(epochs):\n",
        "    loss_per_epoch = 0\n",
        "    for images, labels in train_loader:\n",
        "        images = images.view(images.shape[0], -1)\n",
        "    \n",
        "        # Gradients cleared per batch\n",
        "        optimizer.zero_grad()\n",
        "        \n",
        "        # Pass input to the model\n",
        "        output = model(images)\n",
        "        # Calculate loss after training compared to labels\n",
        "        loss = criterion(output, labels)\n",
        "        \n",
        "        # backpropagation \n",
        "        loss.backward()\n",
        "        \n",
        "        # optimizer step to update the weights\n",
        "        optimizer.step()\n",
        "        \n",
        "        loss_per_epoch += loss.item()\n",
        "    average_loss = loss_per_epoch / len(train_loader)\n",
        "    print(\"Epoch {} - Training loss: {}\".format(epoch, average_loss))"
      ],
      "execution_count": 50,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Epoch 0 - Training loss: 0.749698840479056\n",
            "Epoch 1 - Training loss: 0.22315420204450687\n",
            "Epoch 2 - Training loss: 0.15833292812642952\n",
            "Epoch 3 - Training loss: 0.12278508187967042\n",
            "Epoch 4 - Training loss: 0.10546507445381334\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wn2Wh1AQoAgr"
      },
      "source": [
        "## Model Evaluation\n",
        "\n",
        "Similar to training data loader, we use the validation loader to load batch by batch and run the feed-forward network to get the expected prediction and compared to the label associated with the data point. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zdOIB_sQmX6J",
        "outputId": "2169078c-1f2f-4cf4-d165-4f640d991702"
      },
      "source": [
        "correct_predictions, all_count = 0, 0\n",
        "# enumerate data from the data validation loader (loads a batch at a time)\n",
        "for batch_id, (images,labels) in enumerate(validation_loader):\n",
        "  for i in range(len(labels)):\n",
        "    img = images[i].view(1, 784)\n",
        "    # at prediction stage, only feed-forward calculation is required. \n",
        "    with torch.no_grad():\n",
        "        logps = model(img)\n",
        "\n",
        "    # Output layer of the network uses a LogSoftMax layer\n",
        "    # Hence the probability must be calculated with the exponential values. \n",
        "    # The final layer returns an array of probabilities for each label\n",
        "    # Pick the maximum probability and the corresponding index\n",
        "    # The corresponding index is the predicted label \n",
        "    ps = torch.exp(logps)\n",
        "    probab = list(ps.numpy()[0])\n",
        "    pred_label = probab.index(max(probab))\n",
        "    true_label = labels.numpy()[i]\n",
        "    if(true_label == pred_label):\n",
        "      correct_predictions += 1\n",
        "    all_count += 1\n",
        "\n",
        "print(f\"Model Accuracy {(correct_predictions/all_count) * 100} %\")"
      ],
      "execution_count": 51,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Model Accuracy 96.2 %\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VT0COkaGPlLO"
      },
      "source": [
        "### Reference: \n",
        "\n",
        "1. [Torch NN Sequential](https://pytorch.org/docs/stable/generated/torch.nn.Sequential.html)\n",
        "2. [Handwritten Digit Recognition Using PyTorch — Intro To Neural Networks](https://towardsdatascience.com/handwritten-digit-mnist-pytorch-977b5338e627)\n",
        "3. [MNIST Handwritten Digit Recognition in PyTorch](https://nextjournal.com/gkoehler/pytorch-mnist)\n"
      ]
    }
  ]
}