exercises/03_mnist_classification.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Classifying handwritten digits with a CNN\n",
    "\n",
    "\n",
    "In the previous exercies, we worked with fully-connected neural networks, which are good at handling tabular data, where the inputs and targets are easily presented as vectors.\n",
    "\n",
    "However, in the case of images, or image-like objects, such models are less efficient for reasons we have discussed in the slides. When inputs are images, or image-like data, a more natural choice of model is a convolution neural network—in particularly, a model which uses 2D convolutional layers.\n",
    "\n",
    "Before we start worrying about choosing models, let's first acquaint ourselves with the MNIST data.\n",
    "\n",
    "The first step is to select a directory for the data to live."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mnist_root = \"mnist-data/\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we all set a path this way it will help to maintain consistency throughout this exercise."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task 1: Access the MNIST dataset.\n",
    "\n",
    "PyTorch has a (sort of) sister Python library for dealing with images: [``Torchvision``](https://pytorch.org/vision/stable/index.html) (take a look at the website for a few minutes).\n",
    "\n",
    "In the previous exercies, we used a custom ``Dataset`` object created specifically for this event, but with ``Torchvision`` come several easy-to-use datasets, one of which is the MNIST digits.\n",
    "\n",
    "- Look at the arguments of the MNIST datset: what options do we have?\n",
    "- Instantiate the (training) dataset.\n",
    "- Iterate over it: how are the inputs and targets presented to us?\n",
    "- Plot some images, and set their targets as the title, to make sure the data make sense.\n",
    "\n",
    "Note: this section might be challenging for many people, so we will go through it together first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.datasets import MNIST\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "data_set = MNIST(root=mnist_root, download=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "### Task 2\n",
    "\n",
    "As before, we have a dataset which gives us inputs and targets, but we still need to convert the data to ``torch.Tensors``\n",
    "\n",
    "#### Part (a) — raw data to ``Tensor``s\n",
    "\n",
    "Let's do this basic steps first: supply transforms to map between the raw data and ``torch.Tensors``.\n",
    "\n",
    "\n",
    "\n",
    "#### Part (b) — fun with data augmentation\n",
    "\n",
    "Additionally, with image data there are some other consideration we might make:\n",
    "- Are CNNs rotationally invariant?\n",
    "  - If we want our model to work on images which are not of a regular orientation, we must use random rotations as a form of augmentation.\n",
    "- If we train a model on purely black-and-white images, how will it fare on more colourful data?\n",
    "- Go to ``Torchvision``'s [transforms](https://pytorch.org/vision/stable/transforms.html) and look at the available forms of images augmentation.\n",
    "  - Take a few minutes to pick ones you think might be relevant.\n",
    "  - Let's discuss and choose some:\n",
    "    - ...\n",
    "    - ...\n",
    "    - ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.transforms import ToTensor, Compose\n",
    "from torch import Tensor, tensor\n",
    "\n",
    "\n",
    "def get_img_tfms(training: bool) -> Compose:\n",
    "    \"\"\"Return a composition of image transforms.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    training : bool\n",
    "        Are we training? If ``False``, we are validating.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    Compose\n",
    "        A composition of image transforms.\n",
    "\n",
    "    \"\"\"\n",
    "    # Convert the image from being a 2D tensor to have 3 values for each pixel (RGB).\n",
    "    # Since the three values are the same it remains the same grayscale image\n",
    "\n",
    "    tfm_list: List[Callable] = [ToTensor(), lambda x: x.repeat(3, 1, 1)]\n",
    "\n",
    "    if training is True:\n",
    "        # tfm_list.append(???)\n",
    "        pass\n",
    "\n",
    "    return Compose(tfm_list)\n",
    "\n",
    "\n",
    "def target_as_tensor(target_idx: int) -> Tensor:\n",
    "    \"\"\"Express the target as a ``torch.Tensor``.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    target_idx : int\n",
    "        The ground truth (i.e. the number shown in the image).\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    Tensor\n",
    "        The target as a one-hot-encoded vector.\n",
    "\n",
    "    \"\"\"\n",
    "    return  # ???\n",
    "\n",
    "\n",
    "train_set = MNIST(\n",
    "    mnist_root,\n",
    "    transform=get_img_tfms(training=True),\n",
    "    target_transform=target_as_tensor,\n",
    "    train=True,\n",
    ")\n",
    "\n",
    "\n",
    "valid_set = MNIST(\n",
    "    mnist_root,\n",
    "    transform=get_img_tfms(training=False),\n",
    "    target_transform=target_as_tensor,\n",
    "    train=False,\n",
    ")\n",
    "\n",
    "figure, axes = plt.subplots(5, 5, figsize=(15, 15))\n",
    "\n",
    "for idx, axis in enumerate(axes.ravel()):\n",
    "    image, target = train_set[idx]\n",
    "\n",
    "    axis.imshow(image.permute(1, 2, 0), cmap=\"inferno\")\n",
    "    axis.set_title(target.argmax())\n",
    "    axis.set_axis_off()\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task 3: ``Dataset`` $\\to$ ``DataLoader``\n",
    "\n",
    "As before, wrap the ``Dataset``s in ``DataLoader``s."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "\n",
    "batch_size = 64"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task 4: Choose a model architecture\n",
    "\n",
    "Torchvision provides a collection of models, [here](https://pytorch.org/vision/stable/models.html).\n",
    "\n",
    "Since we are all using laptops, many of which don't have CUDA-enabled GPUs, we will select a modest neural network than won't melt any of our hardware. Such a network, designed for mobile phones, is [``MOBILENET``](https://pytorch.org/vision/stable/models/mobilenetv3.html).\n",
    "\n",
    "#### Task 4 (a): instantiate the small version of ``MOBILENET``, and print it out.\n",
    "\n",
    "Note:\n",
    "- Torchvision's models can optionally be endowed with pretrained weights (from corresponding models pretrain on the ImageNet dataset).\n",
    "- We can (optionally) supply these weights.\n",
    "  - Using weights from a model, trained on one problem, as an initial condition in another problem, is called transfer learning.\n",
    "  - Why do you think this might be advantageous, even in disparate problems?\n",
    "    - ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.models import mobilenet_v3_small\n",
    "from torchvision.models import MobileNet_V3_Small_Weights\n",
    "\n",
    "model = mobilenet_v3_small(weights=MobileNet_V3_Small_Weights.DEFAULT)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Task 4 (b): overloading the final layer\n",
    "\n",
    "Uh oh, we've hit a problem.\n",
    "\n",
    "- Look at the final linear layer.\n",
    "  - How many output classes are there?\n",
    "  - How many do we need?\n",
    "- We need to \"overload\" the final layer to produce the correct number of output features for our problem. Fortunately this is easy.\n",
    "- Uncomment the code below, choose the correct number of output features, and print the model again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn import Linear\n",
    "\n",
    "# model.classifier[3] = Linear(model.classifier[3].in_features, ???)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task 5: Set up the remaining PyTorch bits and bobs\n",
    "\n",
    "- We need to choose a loss function appropriate for classification.\n",
    "  - Can you remember what we chose previously?\n",
    "- We need an optimiser, too.\n",
    "  - Remember our friend, Adam?\n",
    "- Instantiate the model and loss function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn import BCELoss\n",
    "from torch.optim import Adam"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task 6: Set the device\n",
    "\n",
    "We could have done this when we created the model, but it's fine to do it here.  \n",
    "- If [CUDA](https://pytorch.org/docs/stable/cuda.html) (NVIDIA GPU) is available we will elect to run on that.\n",
    "- Alternatively, if we are on Apple Silicon we can use the [MPS backend](https://developer.apple.com/metal/pytorch/) for GPU acceleration.\n",
    "- If neither of these are available we will resort to using the (slower) CPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.cuda as cuda\n",
    "import torch.backends.mps as mps\n",
    "\n",
    "if cuda.is_available():\n",
    "    DEVICE = \"cuda\"\n",
    "elif mps.is_available():\n",
    "    DEVICE = \"mps\"\n",
    "else:\n",
    "    DEVICE = \"cpu\"\n",
    "\n",
    "print(DEVICE)\n",
    "\n",
    "# Note: here the return/assignment to ``_`` is just to suppress the print.\n",
    "# The model is moved onto the correct device in-place.\n",
    "_ = model.to(DEVICE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task 7: Writing our training and validation loops\n",
    "\n",
    "As before, we need to write our training and validation loops.\n",
    "\n",
    "- Complete the training loop\n",
    "- Complete the validation loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Dict\n",
    "\n",
    "from torch import no_grad\n",
    "from torch.nn import Module\n",
    "\n",
    "\n",
    "def train_one_epoch(\n",
    "    model: Module,\n",
    "    train_loader: DataLoader,\n",
    "    optimiser: Adam,\n",
    "    loss_func: BCELoss,\n",
    ") -> Dict[str, float]:\n",
    "    \"\"\"Train the model for a single epoch.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    model : Module\n",
    "        A neural network.\n",
    "    train_loader : DataLoader\n",
    "        The ``DataLoader`` for the training set.\n",
    "    optimiser : Adam\n",
    "        The optimiser to update the model's parameters with.\n",
    "    loss_func : BCELoss\n",
    "        Binary-cross-entropy loss function.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    Dict[str, float]\n",
    "        A dictionary of metrics.\n",
    "\n",
    "    \"\"\"\n",
    "    model.train()\n",
    "\n",
    "\n",
    "@no_grad()\n",
    "def validate_one_epoch(\n",
    "    model: Module,\n",
    "    valid_loader: DataLoader,\n",
    "    loss_func: BCELoss,\n",
    ") -> Dict[str, float]:\n",
    "    \"\"\"Train the model for a single epoch.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    model : Module\n",
    "        A neural network.\n",
    "    valid_loader : DataLoader\n",
    "        The ``DataLoader`` for the validation set.\n",
    "    loss_func : BCELoss\n",
    "        Binary-cross-entropy loss function.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    Dict[str, float]\n",
    "        A dictionary of metrics.\n",
    "\n",
    "    \"\"\"\n",
    "    model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task 8: Training and extracting metrics\n",
    "\n",
    "- Now we can train our model for a specified number of epochs.\n",
    "  - During each epoch the model \"sees\" each training item once.\n",
    "- Append the training and validation metrics to a list.\n",
    "- Turn them into a ``pandas.DataFrame``\n",
    "  - Note: You can turn a ``List[Dict[str, float]]``, say ``my_list`` into a ``DataFrame`` with ``DataFrame(my_list)``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 5\n",
    "\n",
    "for _ in range(epochs):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task 9: Plotting metrics\n",
    "\n",
    " - Let plots the training and validation metrics together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# ..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task 10: Visualising some predictions\n",
    "\n",
    "Let's pick some random validation items, predict on them, and visualise the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "valid_inds = [8, 16, 24, 32, 40]\n",
    "\n",
    "# Stick the model on the cpu\n",
    "_ = model.to(\"cpu\")\n",
    "_ = model.eval()\n",
    "\n",
    "# figure, axes = plt.subplots(1, len(valid_inds), figsize=(2 * len(valid_inds), 2))\n",
    "\n",
    "# for idx, axis in zip(valid_inds, axes.ravel()):\n",
    "#     img_tensor, target = valid_set[idx]\n",
    "\n",
    "#     with no_grad():\n",
    "#         pred = model(img_tensor.unsqueeze(0)).softmax(dim=1).argmax(dim=1).item()\n",
    "\n",
    "#     axis.imshow(img_tensor.permute(1, 2, 0))\n",
    "#     axis.set_axis_off()\n",
    "#     axis.set_title(f\"pred: {pred}, act: {target.argmax()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}