PyTorch_MNIST_tutorial.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "gpuType": "T4",
      "name": "PyTorch MNIST tutorial",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/gileshall/ML-Biology-Notebooks/blob/main/PyTorch_MNIST_tutorial.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Goal: Classify hand-written digits from the MNIST dataset.\n",
        "\n",
        "The MNIST dataset is a collection of grayscale images of handwritten digits (0-9), each 28x28 pixels. The goal is to train a neural network to accurately identify the digit represented in each image. This problem is fundamental to computer vision and serves as a classic example for beginners to understand basic deep learning concepts like image preprocessing, neural network architecture, and model training.\n",
        "\n",
        "The MNIST dataset was created by combining two datasets from the National Institute of Standards and Technology (NIST). It was designed as a simple, accessible benchmark for evaluating image recognition models. It quickly became the \"Hello World\" of deep learning, providing a foundational dataset that is easy to understand but effective for demonstrating neural networks. Despite its simplicity, MNIST played a key role in advancing research and continues to be a popular tool for demonstrating machine learning techniques.\n",
        "\n",
        "The solution involves creating a simple feed-forward neural network to classify MNIST digits implemented using PyTorch. It has several fully connected layers that learn to recognize patterns in the digits. By feeding the network labeled images during training, it learns to distinguish between the digits through optimization of its parameters.\n",
        "\n",
        "This is a good toy example because it encapsulates many core components of deep learning in a simple way."
      ],
      "metadata": {
        "id": "0Dp2pYURUVr1"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xtvUPV53hsOG"
      },
      "outputs": [],
      "source": [
        "!pip install -q ipycanvas orjson\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torch.nn.functional as F\n",
        "from torchvision import datasets, transforms\n",
        "from torch.utils.data import DataLoader\n",
        "import PIL\n",
        "import numpy as np\n",
        "from google.colab import output\n",
        "output.enable_custom_widget_manager()\n",
        "\n",
        "from ipywidgets import Image\n",
        "from ipywidgets import ColorPicker, IntSlider, link, AppLayout, HBox\n",
        "from ipycanvas import Canvas, hold_canvas"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Download and load MNIST hand written digit dataset\n",
        "\n",
        "transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])\n",
        "train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)\n",
        "test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)\n"
      ],
      "metadata": {
        "id": "QI7Auu-kS4Cn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Visualization of the training data\n",
        "\n",
        "for i_exp in range(10):\n",
        "    img = train_dataset[i_exp][0]\n",
        "    img = transforms.ToPILImage()(img)\n",
        "    display(img)"
      ],
      "metadata": {
        "id": "8_czvG_5vxSH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Defining the Neural Network Model\n",
        "\n",
        "This is how to define a custom neural network containing three fully connected layers. The input layer takes 28 x 28 pixels from the flattened MNIST images, followed by two hidden layers with 128 and 64 nodes respectively. The output layer has 10 nodes, each representing one of the digit classes (0-9). The forward propagation function shows how the data moves through the network.\n",
        "\n",
        "Model Architecture:\n",
        "\n",
        "# Neural Network Model for MNIST Handwritten Digit Classification\n",
        "\n",
        "This neural network is a **fully connected (feedforward) neural network** implemented in PyTorch. It aims to classify handwritten digits (0–9) from the **MNIST dataset**, which consists of 28x28 grayscale images.\n",
        "\n",
        "## Model Architecture\n",
        "\n",
        "- **Input Layer:**\n",
        "  - Each image in the MNIST dataset is of size **28x28 pixels**.\n",
        "  - To feed this into the model, it is flattened into a **1D vector of size 784 (28 × 28)** using `x.view(-1, 28 * 28)`.\n",
        "\n",
        "- **First Hidden Layer (fc1):**\n",
        "  - Takes the **784-element vector** from the input layer as input.\n",
        "  - Contains **128 neurons**.\n",
        "  - Applies the **ReLU (Rectified Linear Unit)** activation function.\n",
        "\n",
        "- **Second Hidden Layer (fc2):**\n",
        "  - Takes the **128-element vector** from the first hidden layer as input.\n",
        "  - Contains **64 neurons**.\n",
        "  - Uses the **ReLU activation function** again.\n",
        "\n",
        "- **Output Layer (fc3):**\n",
        "  - Takes the **64-element vector** from the second hidden layer as input.\n",
        "  - Outputs a **10-element vector**, where each element corresponds to one of the **10 classes (digits 0-9)**.\n",
        "\n",
        "Here is a visualization of a fully connected network.\n",
        "\n",
        "![A fully connected network](http://www.gabormelli.com/RKB/images/4/44/2layersFCNN.png \"A fully connected net\")\n",
        "\n",
        "Here is a visualization of a ReLU activation function:\n",
        "\n",
        "![sigmoid and ReLU](https://miro.medium.com/v2/resize:fit:1400/format:webp/1*aEVZlqkcakVySV6ETqgfEg.png \"Sigmoid and ReLU\")"
      ],
      "metadata": {
        "id": "7Z4BLaPAVDwG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Device configuration: Determines whether the model will run on GPU or CPU.\n",
        "# If a CUDA-compatible GPU is available, it will use 'cuda' for faster computation;\n",
        "# otherwise, it will fall back to 'cpu'.\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "\n",
        "# Neural Network Model\n",
        "class NeuralNet(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(NeuralNet, self).__init__()\n",
        "        self.fc1 = nn.Linear(28 * 28, 128)\n",
        "        self.fc2 = nn.Linear(128, 64)\n",
        "        self.fc3 = nn.Linear(64, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x.view(-1, 28 * 28)\n",
        "        x = F.relu(self.fc1(x))\n",
        "        x = F.relu(self.fc2(x))\n",
        "        x = self.fc3(x)\n",
        "        return x\n",
        "\n",
        "# Instantiate the new model\n",
        "model = NeuralNet().to(device)"
      ],
      "metadata": {
        "id": "OC4UJIx0Rj0h"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Hyperparameters: These are key parameters that control the learning process of the model.\n",
        "# Adjusting them can significantly affect the model's performance.\n",
        "\n",
        "# Number of samples processed before the model updates its weights.\n",
        "batch_size = 32\n",
        "\n",
        "# Step size used by the optimizer to update model parameters.\n",
        "learning_rate = 0.01\n",
        "\n",
        "# Total number of complete passes through the entire dataset.\n",
        "epochs = 15\n",
        "\n",
        "# build the data loaders\n",
        "train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)\n",
        "test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)\n",
        "\n",
        "# Loss and optimizer\n",
        "loss_function = nn.CrossEntropyLoss()\n",
        "optimizer = optim.SGD(model.parameters(), lr=learning_rate)\n",
        "\n",
        "# Training loop\n",
        "#\n",
        "# Mode\t    Dropout\t    BatchNorm\t    Gradients\n",
        "# --------------------------------------------------\n",
        "# train()\tActive\t    Batch-wise\t    Computed\n",
        "# eval()\tInactive\tRunning stats\tNot computed\n",
        "\n",
        "for epoch in range(epochs):\n",
        "    model.train()\n",
        "    for batch_idx, (data, target) in enumerate(train_loader):\n",
        "        data, target = data.to(device), target.to(device)\n",
        "\n",
        "        # Forward pass\n",
        "        outputs = model(data)\n",
        "        loss = loss_function(outputs, target)\n",
        "\n",
        "        # Backward pass and optimization\n",
        "        optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "    print(f'Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}')\n"
      ],
      "metadata": {
        "id": "qIb_Re_0lrbt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Test the model\n",
        "model.eval()\n",
        "correct = 0\n",
        "total = 0\n",
        "with torch.no_grad():\n",
        "    for data, target in test_loader:\n",
        "        data, target = data.to(device), target.to(device)\n",
        "        outputs = model(data)\n",
        "        _, predicted = torch.max(outputs.data, 1)\n",
        "        total += target.size(0)\n",
        "        correct += (predicted == target).sum().item()\n",
        "\n",
        "print(f'Accuracy of the model on the test images: {100 * correct / total:.2f}%')\n"
      ],
      "metadata": {
        "id": "-9H3mkrSbp11"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "width = 255\n",
        "height = 255\n",
        "\n",
        "canvas = Canvas(width=width, height=height, sync_image_data=True)\n",
        "\n",
        "drawing = False\n",
        "position = None\n",
        "img_data = [None]\n",
        "\n",
        "def recrop(img):\n",
        "    bbox = img.getbbox()\n",
        "    img = img.crop(bbox)\n",
        "    sq_len = max(img.size) + 32\n",
        "    left_offset = (sq_len - img.size[0]) // 2\n",
        "    top_offset = (sq_len - img.size[1]) // 2\n",
        "    pad_img = PIL.Image.new(img.mode, (sq_len, sq_len), color=(0,))\n",
        "    pad_img.paste(img, (left_offset, top_offset))\n",
        "    return pad_img\n",
        "\n",
        "def predict():\n",
        "    image = PIL.Image.fromarray(canvas.get_image_data())\n",
        "    image = image.convert('L')\n",
        "    image = recrop(image)\n",
        "    # Resize to 28x28 and transform\n",
        "    transform = transforms.Compose([\n",
        "        transforms.ToTensor(),\n",
        "        transforms.Resize((28, 28)),\n",
        "        transforms.Normalize((0.5,), (0.5,))\n",
        "    ])\n",
        "    image = transform(image).unsqueeze(0).to(device)\n",
        "    # show what's going into the network\n",
        "    display(transforms.ToPILImage()(image[0]))\n",
        "\n",
        "    # Predict\n",
        "    model.eval()\n",
        "    with torch.no_grad():\n",
        "        output = model(image)\n",
        "\n",
        "    probabilities = F.softmax(output.data, dim=1)\n",
        "    _, predicted = torch.max(output.data, 1)\n",
        "    probabilities = [f'[{i}: {p * 100:.2f}%]' for (i, p) in enumerate(probabilities[0].cpu().numpy())]\n",
        "    print(f'Probabilities: {\" \".join(probabilities)}]')\n",
        "    print(f'Predicted: {predicted.item()}')\n",
        "\n",
        "def init_canvas():\n",
        "    canvas.fill_style = '#000000'\n",
        "    canvas.fill_rect(0, 0, canvas.width, canvas.height)\n",
        "    canvas.fill_style = '#FFFFFF'\n",
        "    canvas.stroke_style = \"#FFFFFF\"\n",
        "    canvas.line_width = 8\n",
        "\n",
        "def on_mouse_down(x, y):\n",
        "    global drawing\n",
        "    global position\n",
        "\n",
        "    drawing = True\n",
        "    position = (x, y)\n",
        "\n",
        "def on_mouse_move(x, y):\n",
        "    global drawing\n",
        "    global position\n",
        "\n",
        "    if not drawing:\n",
        "        return\n",
        "\n",
        "    with hold_canvas():\n",
        "        canvas.fill_circle(x, y, 7)\n",
        "\n",
        "def on_mouse_up(x, y):\n",
        "    global drawing\n",
        "    global position\n",
        "\n",
        "    drawing = False\n",
        "\n",
        "    with hold_canvas():\n",
        "        canvas.fill_circle(x, y, 7)\n",
        "\n",
        "def on_keyboard_event(key, shift_key, ctrl_key, meta_key):\n",
        "    if key.lower() == 'c':\n",
        "        canvas.clear()\n",
        "        init_canvas()\n",
        "    elif key.lower() == 'p':\n",
        "        predict()\n",
        "\n",
        "canvas.on_key_down(on_keyboard_event)\n",
        "canvas.on_mouse_down(on_mouse_down)\n",
        "canvas.on_mouse_move(on_mouse_move)\n",
        "canvas.on_mouse_up(on_mouse_up)\n",
        "\n",
        "init_canvas()\n",
        "print(\"Draw a digit with your mouse.\\nPress 'p' to predict.\\nPress 'c' to clear.\")\n",
        "canvas"
      ],
      "metadata": {
        "id": "r-wNSEHCnYq3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Exercise: Try a Convolutional Neural Network (CNN)\n",
        "\n",
        "Here is an alternative architecture for you to try. This model uses **Convolutional Neural Networks (CNNs)**, which are especially effective for image classification tasks like recognizing handwritten digits from the MNIST dataset.\n",
        "\n",
        "In this exercise, you can:\n",
        "- Implement this CNN and compare its performance with the fully connected network.\n",
        "- Experiment with the architecture (e.g., change the number of filters, dropout rates, or add more layers).\n",
        "- Train the model on the MNIST dataset and observe how CNNs handle image data more efficiently.\n",
        "\n",
        "To use this model, execute this cell and scroll back up to the training cell.  Comment out the line that creates the NeuralNet model and uncomment the line that creates the CNN model."
      ],
      "metadata": {
        "id": "u4OvL6QDUTRH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define the Convolutional Neural Network\n",
        "class CNN(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(CNN, self).__init__()\n",
        "        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, padding=1)\n",
        "        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1)\n",
        "        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)\n",
        "        self.dropout = nn.Dropout(0.25)\n",
        "        self.fc1 = nn.Linear(64 * 7 * 7, 128)\n",
        "        self.fc2 = nn.Linear(128, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.pool(torch.relu(self.conv1(x)))\n",
        "        x = self.pool(torch.relu(self.conv2(x)))\n",
        "        x = self.dropout(x)\n",
        "        x = x.view(-1, 64 * 7 * 7)  # Flatten the tensor\n",
        "        x = torch.relu(self.fc1(x))\n",
        "        x = self.dropout(x)\n",
        "        x = self.fc2(x)\n",
        "        return x\n",
        "\n",
        "\n",
        "# Alternative model (see below)\n",
        "model = CNN().to(device)"
      ],
      "metadata": {
        "id": "3iiLWjh_eL-H"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}