MLT-2/MLT2_MNIST.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2ce7846b-d162-4f0e-b5f8-657b57b9475e",
   "metadata": {},
   "source": [
    "<img src=\"https://cellstrat2.s3.amazonaws.com/PlatformAssets/bluewhitelogo.svg\" alt=\"drawing\" width=\"200\"/>\n",
    "\n",
    "# ML Tuesdays - Session 2\n",
    "## Deep Learning Track\n",
    "### MNIST Digit Classification Exercise\n",
    "\n",
    "Train a simple neural network in PyTorch to classify handwritten digits using the MNIST dataset.\n",
    "\n",
    "### Guidelines\n",
    "1. Use the `PyTorch 1.9` kernel in CellStrat Hub.\n",
    "2. The notebook has been split into multiple steps.\n",
    "3. Refer the [PyTorch Quickstart Tutorial](https://pytorch.org/tutorials/beginner/basics/quickstart_tutorial.html) as this exercise is based on that.\n",
    "4. Refer the official [PyTorch Documentation](https://pytorch.org/docs/stable/index.html)\n",
    "5. Make use of the docstrings of the functions and classes using the `shift+tab` shortcut key.\n",
    "6. Refer the internet for any extra help needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0acf831-bde9-4874-a3fc-3d964ad46d48",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets\n",
    "from torchvision.transforms import ToTensor"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "235e73e0-adfc-4706-a173-d4fcba4fb847",
   "metadata": {},
   "source": [
    "## Load Data\n",
    "\n",
    "Load the MNIST dataset using torchvision's datasets and add a transform which converts the images to PyTorch Tensors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98ba6b47-8e06-4f0b-a0be-72f3511bad51",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "f5a908f0-be86-4fbb-a254-f3d498f50326",
   "metadata": {},
   "source": [
    "## Create DataLoader\n",
    "Create a dataloader for the dataset and with a batch size of 96."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7508c404-2081-461b-9a42-4655a003c91b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "eeb121ca-bd25-4ec1-8991-45f57ada7ee2",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Build Model\n",
    "\n",
    "Create a simple neural network which has 2 hidden layers atleast with 256 and 512 neurons in each. Try with different network configurations and see if the performance improves after training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2f892-9148-4eb3-9fff-a3800d29729b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get cpu or gpu device automatically for training.\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(f\"Using {device} device\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c976118-c725-4367-ae55-735d646fc13f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NeuralNetwork(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(NeuralNetwork, self).__init__()\n",
    "        # TODO\n",
    "        pass # remove this pass after you complete this\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # TODO\n",
    "        pass # remove this pass after you complete this\n",
    "\n",
    "model = NeuralNetwork().to(device)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3446603-e1c3-43d1-a607-5012bfc19508",
   "metadata": {},
   "source": [
    "## Training\n",
    "\n",
    "1. Define a `CrossEntropyLoss()` loss function (also called _criterion_)\n",
    "2. Create a `Adam` optimizer with learning rate `0.005`\n",
    "3. Write the training and evaluation loop for the train and test sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e9def2e-daf7-4ae4-a5c7-28d5500cc20e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO\n",
    "criterion = None\n",
    "optimizer = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85af0647-6208-46f7-a08c-eee9cefad924",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(dataloader, model, loss_fn, optimizer):\n",
    "    size = len(dataloader.dataset)\n",
    "    \n",
    "    # TODO: set the model to train mode\n",
    "    \n",
    "    for batch, (X, y) in enumerate(dataloader):\n",
    "        # TODO: Move the inputs to device\n",
    "        \n",
    "        # TODO: Model forward pass\n",
    "        \n",
    "        # TODO:  Model forward pass\n",
    "\n",
    "        # TODO: Backpropagation\n",
    "\n",
    "        if batch % 100 == 0:\n",
    "            loss, current = loss.item(), batch * len(X)\n",
    "            print(f\"Train Loss: {loss:>7f}  [{current:>5d}/{size:>5d}]\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de073ceb-d0f1-42aa-85b0-842e2994fd9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(dataloader, model, loss_fn):\n",
    "    size = len(dataloader.dataset)\n",
    "    num_batches = len(dataloader)\n",
    "    \n",
    "    # TODO: set the model to eval mode\n",
    "    \n",
    "    test_loss, correct = 0, 0\n",
    "    with torch.no_grad():\n",
    "        for X, y in dataloader:\n",
    "            pass\n",
    "            # TODO: Move the inputs to device\n",
    "            \n",
    "            # TODO: Model forward pass\n",
    "            \n",
    "            # TODO: Compute the loss and add it to the test_loss variable\n",
    "            \n",
    "            # TODO: Calculate the number of correct predictions and add it to correct variable\n",
    "            \n",
    "    # Calculate Mean loss and accuracy\n",
    "    test_loss /= num_batches\n",
    "    correct /= size\n",
    "    print(f\"Test Error: \\n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5731352-cda1-4626-aab6-1f5ad350dba5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train for 5 epochs\n",
    "epochs = 5\n",
    "\n",
    "for t in range(epochs):\n",
    "    print(f\"Epoch {t+1}\\n-------------------------------\")\n",
    "    train(train_dataloader, model, criterion, optimizer)\n",
    "    test(test_dataloader, model, criterion)\n",
    "    \n",
    "print(\"Training Complete!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0ce9bb5-b953-4a8b-b330-6c0f0bfffaf0",
   "metadata": {},
   "source": [
    "## Saving Model\n",
    "\n",
    "Save the model trained as `mnist_model.pth`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21dd46d7-0fee-4216-90e9-56b8d806e4e3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "43f6f0b0-aa0a-4db0-8a7f-7d60a15ae28a",
   "metadata": {},
   "source": [
    "## Loading Model and Making Predictions\n",
    "\n",
    "1. Create and Load the model using the saved weights\n",
    "2. Make a prediction function so any single image (ex: from the the test data) can be used to get the final predicted class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe044310-7159-4bd0-b9c2-5a5d50eb7de8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Create and Load the model using the saved weights\n",
    "model = NeuralNetwork()\n",
    "model.load_state_dict(torch.load(\"mnist_model.pth\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ab83507-7c6b-4d06-a167-ad069615187d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(model, x):\n",
    "    '''Returns the predicted integer\n",
    "    Args:\n",
    "        model (nn.Module): The loaded model\n",
    "        x (torch.Tensor): A single test image\n",
    "    '''\n",
    "    # set the model in inference mode\n",
    "   \n",
    "        # forward pass on the model\n",
    "        \n",
    "        \n",
    "        # use argmax to get the prediction class integer\n",
    "        \n",
    "        \n",
    "        # return the prediction\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ced4cf3-dedb-4d42-817a-5e196ccb6614",
   "metadata": {},
   "outputs": [],
   "source": [
    "# sample test input\n",
    "x, y = test_data[0]\n",
    "print(x.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3aa49f72-67b3-46d8-9f3b-800b10ebfd63",
   "metadata": {},
   "outputs": [],
   "source": [
    "# make the prediction\n",
    "pred = predict(model, x)\n",
    "print(f'Predicted: {pred}, Actual: {y}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4c69ed7-f8f0-4962-b164-3f4f8786330a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "PyTorch 1.9",
   "language": "python",
   "name": "pytorch"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}