5.1 - Optimization Tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimization activity\n",
    "In this activity, you will write your own optimization algorithm for a MLP on the MNIST dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch  # Main torch import for torch tensors\n",
    "import torch.nn as nn  # Neural network module for building deep learning models\n",
    "import torch.nn.functional as F  # Functional module, includes activation functions\n",
    "import torch.optim as optim  # Optimization module\n",
    "import torchvision  # Vision / image processing package built on top of torch\n",
    "\n",
    "from matplotlib import pyplot as plt  # Plotting and visualization\n",
    "from sklearn.metrics import accuracy_score  # Computing accuracy metric"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup the problem\n",
    "This is all the same as Week 3 MLP tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Common practice to normalize input data to neural networks (0 mean, unit variance)\n",
    "transform = torchvision.transforms.Compose([\n",
    "    torchvision.transforms.ToTensor(),  # All inputs to PyTorch neural networks must be torch.Tensor\n",
    "    torchvision.transforms.Normalize(mean=0.1307, std=0.3081)  # Subtracts mean and divides by std. Note that the raw data is between [0, 1]\n",
    "])\n",
    "\n",
    "# Download the MNIST data and lazily apply the transformation pipeline\n",
    "train_data = torchvision.datasets.MNIST('./datafiles/', train=True, download=True, transform=transform)\n",
    "test_data = torchvision.datasets.MNIST('./datafiles/', train=False, download=True, transform=transform)\n",
    "\n",
    "# Setup data loaders\n",
    "# Note: Iterating through the dataloader yields batches of (inputs, targets)\n",
    "# where inputs is a torch.Tensor of shape (B, 1, 28, 28) and targets is a torch.Tensor of shape (B,)\n",
    "train_loader = torch.utils.data.DataLoader(train_data, batch_size=64, shuffle=True)\n",
    "test_loader = torch.utils.data.DataLoader(test_data, batch_size=1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(4, 5, figsize=(5, 6))\n",
    "\n",
    "plot_images = []\n",
    "plot_labels = []\n",
    "\n",
    "for i, ax in enumerate(axs.flatten(), start=1000):\n",
    "    (image, label) = test_data[i]\n",
    "\n",
    "    # Save this data for later\n",
    "    plot_images.append(image)\n",
    "    plot_labels.append(label)\n",
    "\n",
    "    # Plot each image\n",
    "    ax.imshow(image.squeeze(), cmap=\"gray\")\n",
    "    ax.set_title(f\"Label: {label}\")\n",
    "    ax.axis(\"off\")\n",
    "plt.show()\n",
    "\n",
    "plot_images = torch.cat(plot_images)  # Combine all the images into a single batch for later\n",
    "\n",
    "print(f\"Each image is a torch.Tensor and has shape {image.shape}.\")\n",
    "print(f\"The labels are the integers 0 to 9, representing the digits.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiLayerPerceptron(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.hidden = nn.Linear(784, 100)\n",
    "        self.output = nn.Linear(100, 10)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Forward pass implementation for the network\n",
    "        \n",
    "        :param x: torch.Tensor of shape (batch, 1, 28, 28), input images\n",
    "\n",
    "        :returns: torch.Tensor of shape (batch, 10), output logits\n",
    "        \"\"\"\n",
    "        x = torch.flatten(x, 1)  # shape (batch, 28*28)\n",
    "        x = self.hidden(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.output(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup a custom optimizer\n",
    "\n",
    "Here, we will use a class with a similar interface to PyTorch's optimizers and implement SGD.\n",
    "\n",
    "Try implementing your own version of another optimizer in the slides (but not Nesterov momentum).\n",
    "\n",
    "Relevant documentation:\n",
    "- [PyTorch optimizers](https://pytorch.org/docs/stable/optim.html)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MySGD:\n",
    "    def __init__(self, params, lr=1e-2):\n",
    "        pass\n",
    "    \n",
    "    def zero_grad(self):\n",
    "        pass\n",
    "\n",
    "    def step(self):\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test one step\n",
    "model = MultiLayerPerceptron()\n",
    "optimizer = MySGD(model.parameters(), lr=1e-1)\n",
    "\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "inputs = torch.randn((4, 1, 28, 28))\n",
    "targets = torch.LongTensor([0, 1, 2, 3])\n",
    "\n",
    "logits = model(inputs)\n",
    "loss = loss_fn(logits, targets)\n",
    "loss.backward()\n",
    "optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, train_loader, loss_fn, optimizer, epoch=-1):\n",
    "    \"\"\"\n",
    "    Trains a model for one epoch (one pass through the entire training data).\n",
    "\n",
    "    :param model: PyTorch model\n",
    "    :param train_loader: PyTorch Dataloader for training data\n",
    "    :param loss_fn: PyTorch loss function\n",
    "    :param optimizer: PyTorch optimizer, initialized with model parameters\n",
    "    :kwarg epoch: Integer epoch to use when printing loss and accuracy\n",
    "    :returns: Accuracy score\n",
    "    \"\"\"\n",
    "    total_loss = 0\n",
    "    all_predictions = []\n",
    "    all_targets = []\n",
    "    loss_history = []\n",
    "\n",
    "    model.train()  # Set model in training mode\n",
    "    for i, (inputs, targets) in enumerate(train_loader):\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(inputs)\n",
    "        loss = loss_fn(outputs, targets)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "\n",
    "        # Track some values to compute statistics\n",
    "        total_loss += loss.item()\n",
    "        preds = torch.argmax(outputs, dim=-1)\n",
    "        all_predictions.extend(preds.tolist())\n",
    "        all_targets.extend(targets.tolist())\n",
    "\n",
    "        # Save loss every 100 batches\n",
    "        if (i % 100 == 0) and (i > 0):\n",
    "            running_loss = total_loss / (i + 1)\n",
    "            loss_history.append(running_loss)\n",
    "            # print(f\"Epoch {epoch + 1}, batch {i + 1}: loss = {running_loss:.2f}\")\n",
    "\n",
    "    acc = accuracy_score(all_targets, all_predictions)\n",
    "    final_loss = total_loss / len(train_loader)\n",
    "    # Print average loss and accuracy\n",
    "    print(f\"Epoch {epoch + 1} done. Average train loss = {final_loss:.2f}, average train accuracy = {acc * 100:.3f}%\")\n",
    "    return acc, final_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(model, test_loader, loss_fn, epoch=-1):\n",
    "    \"\"\"\n",
    "    Tests a model for one epoch of test data.\n",
    "\n",
    "    Note:\n",
    "        In testing and evaluation, we do not perform gradient descent optimization, so steps 2, 5, and 6 are not needed.\n",
    "        For performance, we also tell torch not to track gradients by using the `with torch.no_grad()` context.\n",
    "\n",
    "    :param model: PyTorch model\n",
    "    :param test_loader: PyTorch Dataloader for test data\n",
    "    :param loss_fn: PyTorch loss function\n",
    "    :kwarg epoch: Integer epoch to use when printing loss and accuracy\n",
    "\n",
    "    :returns: Accuracy score\n",
    "    \"\"\"\n",
    "    total_loss = 0\n",
    "    all_predictions = []\n",
    "    all_targets = []\n",
    "\n",
    "    model.eval()  # Set model in evaluation mode\n",
    "    for i, (inputs, targets) in enumerate(test_loader):\n",
    "        with torch.no_grad():\n",
    "            outputs = model(inputs)\n",
    "            loss = loss_fn(outputs, targets)\n",
    "\n",
    "            # Track some values to compute statistics\n",
    "            total_loss += loss.item()\n",
    "            preds = torch.argmax(outputs, dim=-1)\n",
    "            all_predictions.extend(preds.tolist())\n",
    "            all_targets.extend(targets.tolist())\n",
    "\n",
    "    acc = accuracy_score(all_targets, all_predictions)\n",
    "\n",
    "    # Print average loss and accuracy\n",
    "    print(f\"Epoch {epoch + 1} done. Average test loss = {total_loss / len(test_loader):.2f}, average test accuracy = {acc * 100:.3f}%\")\n",
    "    return acc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test the implementation\n",
    "Train a model for 5 epochs with both custom and PyTorch optimizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LEARNING_RATE = 1e-1\n",
    "NUM_EPOCHS = 5\n",
    "torch.manual_seed(0)\n",
    "model = MultiLayerPerceptron()\n",
    "optimizer = MySGD(model.parameters(), lr=LEARNING_RATE)\n",
    "loss_fn = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_metrics = []\n",
    "test_metrics = []\n",
    "train_losses_custom = []\n",
    "for epoch in range(NUM_EPOCHS):\n",
    "    train_acc, train_loss = train(model, train_loader, loss_fn, optimizer, epoch)\n",
    "    train_losses_custom.append(train_loss)\n",
    "    test_acc = test(model, test_loader, loss_fn, epoch)\n",
    "\n",
    "    train_metrics.append(train_acc)\n",
    "    test_metrics.append(test_acc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "model = MultiLayerPerceptron()\n",
    "optimizer = optim.SGD(model.parameters(), lr=LEARNING_RATE)\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "\n",
    "train_metrics = []\n",
    "test_metrics = []\n",
    "train_losses_torch = []\n",
    "for epoch in range(NUM_EPOCHS):\n",
    "    train_acc, train_loss = train(model, train_loader, loss_fn, optimizer, epoch)\n",
    "    train_losses_torch.append(train_loss)\n",
    "    test_acc = test(model, test_loader, loss_fn, epoch)\n",
    "\n",
    "    train_metrics.append(train_acc)\n",
    "    test_metrics.append(test_acc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visually compare the model predictions\n",
    "\n",
    "We will lastly see the trained model's predictions on the 20 examples we visualized in the beginning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(train_losses_torch, c=\"r\", label=\"Torch optimizer\")\n",
    "plt.plot(train_losses_custom, c=\"b\", label=\"Custom optimizer\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "datasci",
   "language": "python",
   "name": "python3"
  },
  "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.10.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}