{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "Copyright (c) 2020 Graphcore Ltd. All rights reserved."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "*Notebook autogenerated from walkthrough.py on 27-Sep-2022*"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Introduction to PopTorch: Running a Simple Model\n",
    "\n",
    "This tutorial covers the basics of model making in PyTorch, using\n",
    "`torch.nn.Module`, and the specific methods to convert a PyTorch model to\n",
    "a PopTorch model so that it can be run on a Graphcore IPU."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Running on Paperspace\n",
    "\n",
    "The Paperspace environment lets you run this notebook with no set up. To improve your experience we preload datasets and pre-install packages, this can take a few minutes, if you experience errors immediately after starting a session please try restarting the kernel before contacting support. If a problem persists or you want to give us feedback on the content of this notebook, please reach out to through our community of developers using our [slack channel](https://www.graphcore.ai/join-community) or raise a [GitHub issue](https://github.com/gradient-ai/Graphcore-Pytorch/issues)."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "%pip install -q -r requirements.txt"
   ],
   "outputs": [],
   "metadata": {
    "tags": [
     "sst_ignore_md",
     "sst_ignore_code_only"
    ]
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "265869fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "from examples_utils import notebook_logging\n",
    "%load_ext gc_logger"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "## What is PopTorch?\n",
    "\n",
    "PopTorch is a set of extensions for PyTorch to enable PyTorch models to run\n",
    "on Graphcore's IPU hardware.\n",
    "\n",
    "PopTorch supports both inference and training. To run a model on the IPU you\n",
    "wrap your existing PyTorch model in either a PopTorch inference wrapper or\n",
    "a PopTorch training wrapper. You can provide further annotations to partition\n",
    "the model across multiple IPUs.\n",
    "\n",
    "You can wrap individual layers in an IPU helper to designate which IPU they\n",
    "should go on. Using your annotations, PopTorch will use\n",
    "[PopART](https://docs.graphcore.ai/projects/popart-user-guide) to parallelise\n",
    "the model over the given number of IPUs. Additional parallelism\n",
    "can be expressed via a replication factor which enables you to\n",
    "data-parallelise the model over more IPUs.\n",
    "\n",
    "Under the hood PopTorch uses\n",
    "[TorchScript](https://pytorch.org/docs/stable/jit.html), an intermediate\n",
    "representation (IR) of a PyTorch model, using the `torch.jit.trace` API. That\n",
    "means it inherits the constraints of that API. These include:\n",
    "\n",
    "- Inputs must be Torch tensors or tuples/lists containing Torch tensors\n",
    "- `None` can be used as a default value for a parameter but cannot be explicitly\n",
    "  passed as an input value\n",
    "- Hooks and `.grad` cannot be used to inspect weights and gradients\n",
    "- `torch.jit.trace` cannot handle control flow or shape variations within the\n",
    "  model. That is, the inputs passed at run-time cannot vary the control flow of\n",
    "  the model or the shapes/sizes of results.\n",
    "\n",
    "To learn more about TorchScript and JIT, you can go through the [Introduction\n",
    "to TorchScript tutorial](https://pytorch.org/tutorials/beginner/Intro_to_TorchScript_tutorial.html).\n",
    "\n",
    "PopTorch has been designed to require only a few manual changes to your models\n",
    "in order to run them on the IPU. However, it does have some differences from\n",
    "native PyTorch execution and not all PyTorch operations have been\n",
    "implemented in the backend yet. You can find the list of supported operations\n",
    "[here](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.1.0/supported_ops.html).\n",
    "\n",
    "![Software stack](static/stack.jpg)"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Getting started: training a model on the IPU\n",
    "\n",
    "We will do the following steps in order:\n",
    "\n",
    "1. Load the Fashion-MNIST dataset using `torchvision.datasets` and\n",
    "   `poptorch.DataLoader`.\n",
    "2. Define a deep CNN  and a loss function using the `torch` API.\n",
    "3. Train the model on an IPU using `poptorch.trainingModel`.\n",
    "4. Evaluate the model on the IPU."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Import the packages\n",
    "\n",
    "PopTorch is a separate package from PyTorch, and available\n",
    "in Graphcore's Poplar SDK. Both must thus be imported:"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import torch\n",
    "import poptorch\n",
    "import torchvision\n",
    "import torch.nn as nn\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm import tqdm\n",
    "\n",
    "# Set torch random seed for reproducibility\n",
    "torch.manual_seed(42)"
   ],
   "outputs": [],
   "metadata": {
    "tags": [
     "sst_hide_output"
    ]
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Under the hood, PopTorch uses Graphcore's high-performance\n",
    "machine learning framework PopART. It is therefore necessary\n",
    "to enable PopART and Poplar in your environment."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "> **NOTE**:\n",
    ">\n",
    "> If you forget to enable PopART, you will encounter the following error when\n",
    "> importing: `poptorch`: `ImportError: libpopart.so: cannot open shared object\n",
    "> file: No such file or directory`\n",
    ">\n",
    "> If the error message says something like: `libpopart_compiler.so: undefined\n",
    "> symbol: _ZN6popart7Session3runERNS_7IStepIOE`, it most likely means the\n",
    "> versions of PopART and PopTorch do not match, for example because\n",
    "> you have sourced the PopART `enable.sh` script for a different SDK version to\n",
    "> that of PopTorch. Ensure that you use the same version of the SDK for PopTorch\n",
    "> and PopART."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Load the data\n",
    "\n",
    "We will use the Fashion-MNIST dataset made available by the package\n",
    "`torchvision`. This dataset, from\n",
    "[Zalando](https://github.com/zalandoresearch/fashion-mnist), can be used as a\n",
    "more challenging replacement to the well-known MNIST dataset."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "The dataset consists of 28x28 grayscale images and labels of range `[0, 9]`\n",
    "from 10 classes: T-shirt, trouser, pullover, dress, coat, sandal, shirt,\n",
    "sneaker, bag and ankle boot."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "In order for the images to be usable by PyTorch, we have to convert them to\n",
    "`torch.Tensor` objects. Also, data normalisation improves overall\n",
    "performance. We will apply both operations, conversion and normalisation, to\n",
    "the datasets using `torchvision.transforms` and feed these ops to\n",
    "`torchvision.datasets`:"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "transform = torchvision.transforms.Compose(\n",
    "    [\n",
    "        torchvision.transforms.ToTensor(),\n",
    "        torchvision.transforms.Normalize((0.5,), (0.5,)),\n",
    "    ]\n",
    ")\n",
    "\n",
    "train_dataset = torchvision.datasets.FashionMNIST(\n",
    "    \"~/.torch/datasets\", transform=transform, download=True, train=True\n",
    ")\n",
    "\n",
    "test_dataset = torchvision.datasets.FashionMNIST(\n",
    "    \"~/.torch/datasets\", transform=transform, download=True, train=False\n",
    ")\n",
    "\n",
    "classes = (\n",
    "    \"T-shirt\",\n",
    "    \"Trouser\",\n",
    "    \"Pullover\",\n",
    "    \"Dress\",\n",
    "    \"Coat\",\n",
    "    \"Sandal\",\n",
    "    \"Shirt\",\n",
    "    \"Sneaker\",\n",
    "    \"Bag\",\n",
    "    \"Ankle boot\",\n",
    ")"
   ],
   "outputs": [],
   "metadata": {
    "tags": [
     "sst_hide_output"
    ]
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "With the following method, we can visualise and save a sample of these images\n",
    "and their associated labels:"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "plt.figure(figsize=(30, 15))\n",
    "for i, (image, label) in enumerate(train_dataset):\n",
    "    if i == 15:\n",
    "        break\n",
    "    image = (image / 2 + 0.5).numpy()  # reverse transformation\n",
    "    ax = plt.subplot(5, 5, i + 1)\n",
    "    ax.set_title(classes[label])\n",
    "    plt.imshow(image[0])\n",
    "\n",
    "plt.savefig(\"sample_images.png\")"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### PopTorch DataLoader\n",
    "\n",
    "We can feed batches of data into a PyTorch model by simply passing the input\n",
    "tensors. However, this is unlikely to be efficient and can\n",
    "result in data loading being a bottleneck to the model, slowing down the\n",
    "training process. In order to make data loading easier and more efficient,\n",
    "there's the\n",
    "[`torch.utils.data.DataLoader`](https://pytorch.org/docs/stable/data.html)\n",
    "class, which is an `iterable` over a dataset and which can handle parallel data\n",
    "loading, a sampling strategy, shuffling, etc."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "PopTorch offers an extension of this class with its\n",
    "[`poptorch.DataLoader`](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.1.0/batching.html#poptorch-dataloader)\n",
    "class, specialised for the way the underlying PopART framework handles\n",
    "batching of data. We will use this class later in the tutorial, as soon as we\n",
    "have a model ready for training."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Build the model\n",
    "\n",
    "We will build a simple CNN model for a classification task. To do so, we can\n",
    "simply use PyTorch's API, including `torch.nn.Module`. The difference from\n",
    "what we're used to with pure PyTorch is the _loss computation_, which has to\n",
    "be part of the `forward` function. This is to ensure the loss is computed on\n",
    "the IPU and not on the CPU, and to give us as much flexibility as possible\n",
    "when designing more complex loss functions."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "class ClassificationModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 5, 3)\n",
    "        self.pool = nn.MaxPool2d(2, 2)\n",
    "        self.conv2 = nn.Conv2d(5, 12, 5)\n",
    "        self.norm = nn.GroupNorm(3, 12)\n",
    "        self.fc1 = nn.Linear(972, 100)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.fc2 = nn.Linear(100, 10)\n",
    "        self.log_softmax = nn.LogSoftmax(dim=1)\n",
    "        self.loss = nn.NLLLoss()\n",
    "\n",
    "    def forward(self, x, labels=None):\n",
    "        x = self.pool(self.relu(self.conv1(x)))\n",
    "        x = self.norm(self.relu(self.conv2(x)))\n",
    "        x = torch.flatten(x, start_dim=1)\n",
    "        x = self.relu(self.fc1(x))\n",
    "        x = self.log_softmax(self.fc2(x))\n",
    "        # The model is responsible for the calculation\n",
    "        # of the loss when using an IPU. We do it this way:\n",
    "        if self.training:\n",
    "            return x, self.loss(x, labels)\n",
    "        return x\n",
    "\n",
    "\n",
    "model = ClassificationModel()\n",
    "model.train()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "> **NOTE**: `self.training` is inherited from `torch.nn.Module` which\n",
    "> initialises its value to `True`. Use `model.eval()` to set it to `False` and\n",
    "> `model.train()` to switch it back to `True`."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Prepare training for IPUs\n",
    "\n",
    "The compilation and execution on the IPU can be controlled using\n",
    "`poptorch.Options`. These options are used by PopTorch's wrappers such as\n",
    "`poptorch.DataLoader` and `poptorch.trainingModel`."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "opts = poptorch.Options()\n",
    "\n",
    "train_dataloader = poptorch.DataLoader(\n",
    "    opts, train_dataset, batch_size=16, shuffle=True, num_workers=20\n",
    ")"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Train the model\n",
    "\n",
    "We will need another component in order to train our model: an optimiser.\n",
    "Its role is to apply the computed gradients to the model's weights to optimize\n",
    "(usually, minimize) the loss function using a specific algorithm. PopTorch\n",
    "currently provides classes which inherit from multiple native PyTorch\n",
    "optimisation functions: SGD, Adam, AdamW, LAMB and RMSprop. These optimisers\n",
    "provide several advantages over native PyTorch versions. They embed constant\n",
    "attributes to save performance and memory, and allow you to specify additional\n",
    "parameters such as loss/velocity scaling."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "We will use\n",
    "[SGD](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.1.0/reference.html#poptorch.optim.SGD)\n",
    "as it's a very popular algorithm and is appropriate for this classification\n",
    "task."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "optimizer = poptorch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)"
   ],
   "outputs": [],
   "metadata": {
    "tags": [
     "sst_hide_output"
    ]
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "We now introduce the `poptorch.trainingModel` wrapper, which will handle the\n",
    "training. It takes an instance of `torch.nn.Module`, such as our custom\n",
    "model, an instance of `poptorch.Options` which we have instantiated\n",
    "previously, and an optimizer. This wrapper will trigger the compilation of\n",
    "our model, using TorchScript, and manage its translation to a program the\n",
    "IPU can run. Let's use it."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "poptorch_model = poptorch.trainingModel(model, options=opts, optimizer=optimizer)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Training loop\n",
    "\n",
    "Looping through the training data, running the forward and backward passes,\n",
    "and updating the weights constitute the process we refer to as the \"training\n",
    "loop\". Graphcore's Poplar graph program framework uses several optimisations to accelerate the\n",
    "training loop. Central to this is the desire to minimise interactions between\n",
    "the device (the IPU) and the host (the CPU), allowing the training loop to\n",
    "run on the device independently from the host. To achieve that virtual\n",
    "independence, Poplar creates a _static_ computational graph and data streams\n",
    "which are loaded to the IPU, and then signals the IPU to get started until\n",
    "there's no data left or until the host sends a signal to stop the loop.\n",
    "\n",
    "![High-level overview of what happens](static/loop.jpg)"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "The compilation, which transforms our PyTorch model into a computational\n",
    "graph and our dataloader into data streams, happens at the first call of a\n",
    "`poptorch.trainingModel`. The IPUs to which the graph will be uploaded are\n",
    "selected automatically during this first call, by default. The training loop\n",
    "can then start."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "Once the loop has started, Poplar's main task is to feed the data into the\n",
    "streams and to signal when we are done with the loop. The last step will then\n",
    "be to copy the final graph, meaning the model, back to the CPU - a step that\n",
    "PopTorch manages itself."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "epochs = 5\n",
    "for epoch in tqdm(range(epochs), desc=\"epochs\"):\n",
    "    total_loss = 0.0\n",
    "    for data, labels in tqdm(train_dataloader, desc=\"batches\", leave=False):\n",
    "        output, loss = poptorch_model(data, labels)\n",
    "        total_loss += loss"
   ],
   "outputs": [],
   "metadata": {
    "tags": [
     "sst_hide_output"
    ]
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "The model is now trained! There's no need to retrieve the weights from the\n",
    "device as you would by calling `model.cpu()` with PyTorch. PopTorch has\n",
    "managed that step for us. We can now save and evaluate the model."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Use the same IPU for training and inference\n",
    "\n",
    "After the model has been attached to the IPU and compiled after the first call\n",
    "to the PopTorch model, it can be detached from the device. This allows PopTorch\n",
    "to use a single device for training and inference (described below), rather\n",
    "than using 2 IPUs (one for training and one for inference) when the device\n",
    "is not detached. When using a POD system, detaching from the device will\n",
    "be necessary when using a non-reconfigurable partition."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "poptorch_model.detachFromDevice()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Save the trained model\n",
    "\n",
    "We can simply use PyTorch's API to save a model in a `.pth` file, with the original\n",
    "instance of `ClassificationModel` and not the wrapped model.\n",
    "\n",
    "Do not hesitate to experiment with different models: the model provided in this\n",
    "tutorial is saved in the `static` folder if you need it."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "torch.save(model.state_dict(), \"classifier.pth\")"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Evaluate the model\n",
    "\n",
    "The model can be evaluated on a CPU but it is a good idea to use the IPU -\n",
    "since [IPUs are blazing\n",
    "fast](https://www.graphcore.ai/posts/new-graphcore-ipu-benchmarks)!\n",
    "\n",
    "Evaluating your model on a CPU is slow if the test dataset is large and/or the\n",
    "model is complex.\n",
    "\n",
    "Since we have detached our model from its training device, the device is now\n",
    "free again and we can use it for the evaluation stage.\n",
    "\n",
    "The steps taken below to define the model for evaluation essentially allow it\n",
    "to run in inference mode. Therefore, you can follow the same steps to use the\n",
    "model to make predictions once it has been deployed."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "model = model.eval()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "To evaluate the model on the IPU, we will use the `poptorch.inferenceModel`\n",
    "class, which has a similar API to `poptorch.trainingModel` except that it\n",
    "doesn't need an optimizer, allowing evaluation of the model without calculating\n",
    "gradients."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "poptorch_model_inf = poptorch.inferenceModel(model, options=opts)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "Then we can instantiate a new PopTorch `Dataloader` object as before in order to\n",
    "efficiently batch our test dataset."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "test_dataloader = poptorch.DataLoader(opts, test_dataset, batch_size=32, num_workers=10)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "This short loop over the test dataset is effectively all that is needed to\n",
    "run the model and generate some predictions. When running the model in\n",
    "inference mode, we can stop here and use the predictions as needed."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "For evaluation, we can use the standard classification metrics from `scikit-learn`\n",
    "to understand how well our model is performing. This usually takes a list\n",
    "of labels and a list of predictions as the input, both in the same order.\n",
    "Let's make both lists, and run our model in inference mode."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "predictions, labels = [], []\n",
    "for data, label in test_dataloader:\n",
    "    predictions += poptorch_model_inf(data).data.max(dim=1).indices\n",
    "    labels += label"
   ],
   "outputs": [],
   "metadata": {
    "tags": [
     "sst_hide_output"
    ]
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Release IPU resources again:"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "poptorch_model_inf.detachFromDevice()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "A simple and widely-used performance metric for classification models is the\n",
    "accuracy score, which simply counts how many predictions were right. But this\n",
    "metric alone isn't enough. For example, it doesn't tell us how the model\n",
    "performs with regard to the different classes in our data. We will therefore\n",
    "use another popular metric: a confusion matrix, which indicates how much our\n",
    "model confuses one class for another."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "from sklearn.metrics import accuracy_score, confusion_matrix, ConfusionMatrixDisplay\n",
    "\n",
    "print(f\"Eval accuracy: {100 * accuracy_score(labels, predictions):.2f}%\")\n",
    "cm = confusion_matrix(labels, predictions)\n",
    "cm_plot = ConfusionMatrixDisplay(cm, display_labels=classes).plot(\n",
    "    xticks_rotation=\"vertical\"\n",
    ")"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "As you can see, although we've got an accuracy score of ~89%, the model's\n",
    "performance across the different classes isn't equal. Trousers are very well\n",
    "classified, with more than 96-97% accuracy whereas shirts are harder to\n",
    "classify with less than 60% accuracy, and it seems they often get confused\n",
    "with T-shirts, pullovers and coats. So, some work is still required here to\n",
    "improve your model for all the classes!"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "We can save this visualisation of the confusion matrix. Don't hesitate to\n",
    "experiment: you can then compare your confusion matrix with the\n",
    "[visualisation provided in the `static` folder](static/confusion_matrix.png)."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "cm_plot.figure_.savefig(\"confusion_matrix.png\")"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Using the model on our own images to get predictions\n",
    "\n",
    "Now that we have a trained model which has also been validated, we can now use\n",
    "the model on arbitrary images which are not part of the FashionMNIST dataset. In\n",
    "this section we will use an image downloaded from the internet and predict the\n",
    "class of it."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Running our model on the IPU"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "We can then instantiate our model and load its `state_dict` which we [saved\n",
    "earlier](#save-the-trained-model). We need to make sure that the model is then set to evaluation\n",
    "mode. We can then wrap our model with the PopTorch inference model which will\n",
    "let us run it on the IPU."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "model = ClassificationModel()\n",
    "model.load_state_dict(torch.load(\"classifier.pth\"))\n",
    "model.eval()\n",
    "\n",
    "poptorch_model = poptorch.inferenceModel(model, options=poptorch.Options())"
   ],
   "outputs": [],
   "metadata": {
    "tags": [
     "sst_hide_output"
    ]
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Next we can load our image into our Python program. To do this we will use the\n",
    "Python Imaging Library (PIL) by supplying our image as a path. To get our image\n",
    "ready for evaluation, we need to make sure that the image is in the same\n",
    "'format' as the images which were in the FashionMNIST dataset. This means that\n",
    "we need to resize our image to 28x28 pixels and make it grayscale.\n",
    "\n",
    "Note: We can download images on a remote machine using wget:\n",
    "`wget \"https://bit.ly/3uKBBDS\" -O ./images/trousers.jpg`\n",
    "\n",
    "> Image attribution: Elisabeth Eriksson / Nordiska museet, CC BY 4.0\n",
    "> <https://creativecommons.org/licenses/by/4.0>, via Wikimedia Commons"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "from PIL import Image, ImageOps\n",
    "\n",
    "img = Image.open(\"images/trousers.jpg\").resize((28, 28))\n",
    "img = ImageOps.grayscale(img)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "Another factor we have to consider is that all of the images in the FashionMNIST\n",
    "dataset have a black background. This means that if our image's background is\n",
    "light, then the model will incorrectly classify the image (another limitation\n",
    "due to our simple model). We can get around this by checking the pixel value in\n",
    "one corner of the image and if it's light, then invert the whole image.\n",
    "Afterwards we can transform our image into a tensor which can be fed into the\n",
    "model."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "if img.getpixel((1, 1)) > 200:\n",
    "    img = ImageOps.invert(img)\n",
    "\n",
    "transform = torchvision.transforms.Compose(\n",
    "    [\n",
    "        torchvision.transforms.ToTensor(),\n",
    "        torchvision.transforms.Normalize((0.5,), (0.5,)),\n",
    "    ]\n",
    ")\n",
    "\n",
    "img_tensor = transform(img)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "We can visualise our original image before and after the transformations:"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "f, ax = plt.subplots(1, 2)\n",
    "ax[0].imshow(Image.open(\"images/trousers.jpg\"))\n",
    "ax[0].set_title(\"Original image\")\n",
    "img_transform = torchvision.transforms.ToPILImage()\n",
    "ax[1].imshow((img_transform(img_tensor)).convert(\"RGBA\"))\n",
    "ax[1].set_title(\"Transformed image\")\n",
    "\n",
    "plt.savefig(\"image_before_after.png\")"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "The model expects a four-dimensional tensor with the dimensions corresponding\n",
    "respectively to: batch size, number of colour channels, and the height and width\n",
    "of the image. Since we have a single image represented as a three-dimensional\n",
    "tensor, we expand it to include a batch dimension using the `unsqueeze` function.\n",
    "\n",
    "The output from the model will be a tensor consisting of 10 values (one for\n",
    "each class). The highest value within the tensor will be the piece of clothing\n",
    "which the model predicts our image to be."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "img_tensor = img_tensor.unsqueeze(0)\n",
    "\n",
    "output = poptorch_model(img_tensor)\n",
    "print(output)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now we interpret the output and print the name of the class."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "prediction_idx = int(output.argmax())\n",
    "\n",
    "poptorch_model.detachFromDevice()\n",
    "\n",
    "print(\"IPU predicted class:\", classes[prediction_idx])"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Running our model on the CPU"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "If you wanted to use the model on a different machine, perhaps not the IPU, then\n",
    "you can import the classification model and the tuple of classes from earlier.\n",
    "This would mean that you can run the prediction without having to import\n",
    "PopTorch at all and we would just need the `.pth` model file."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "model = ClassificationModel()\n",
    "model.load_state_dict(torch.load(\"classifier.pth\"))\n",
    "model.eval()"
   ],
   "outputs": [],
   "metadata": {
    "tags": [
     "sst_hide_output"
    ]
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "output = model(img_tensor)\n",
    "print(\"CPU predicted class:\", classes[int(output.argmax())])"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Limitations with our model\n",
    "\n",
    "Due to the small size of the training dataset and low complexity of the\n",
    "model, there are many limitations with our model. The examples need to be\n",
    "carefully selected to match the distribution of the training dataset, otherwise\n",
    "our model may classify incorrectly.\n",
    "\n",
    "The following need to be considered when choosing images for classification:\n",
    "- The orientation of the piece of clothing must be horizontal,\n",
    "- The background of the image must be a solid colour,\n",
    "- The piece of clothing is photographed on its own; it cannot be worn by a person."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Doing more with `poptorch.Options`\n",
    "\n",
    "This class encapsulates the options that PopTorch and PopART will use\n",
    "alongside our model. Some concepts, such as \"batch per iteration\" are\n",
    "specific to the functioning of the IPU, and within this class some\n",
    "calculations are made to reduce risks of errors and make it easier for\n",
    "PyTorch users to use IPUs."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "The list of these options is available in the\n",
    "[documentation](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.1.0/overview.html#options).\n",
    "We introduce here four of these options so you get an idea of what they cover."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### `deviceIterations`\n",
    "\n",
    "Remember the [training loop](#training-loop) we have discussed previously. A device iteration\n",
    "is one cycle of that loop, which runs entirely on the IPU (the device), and\n",
    "which starts with a new batch of data. This option specifies the number of\n",
    "batches that is prepared by the host (CPU) for the IPU. The higher this\n",
    "number, the less the IPU has to interact with the CPU, for example to request\n",
    "and wait for data, so that the IPU can loop faster. However, the user will\n",
    "have to wait for the IPU to complete all iterations before getting the\n",
    "results back. The maximum number of iterations is the total number of batches in your dataset, and\n",
    "the default value is 1."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### `replicationFactor`\n",
    "\n",
    "This is the number of replicas of a model. A replica is a copy of the same\n",
    "model on multiple devices. We use replicas as an implementation of data\n",
    "parallelism, where the same model is served with several batches of data at the\n",
    "same time but on different devices, so that the gradients can be pooled. To\n",
    "achieve the same behaviour in pure PyTorch, you'd wrap your model with\n",
    "`torch.nn.DataParallel`, but with PopTorch, this is an option. Of course, each\n",
    "replica requires one IPU. So, if the `replicationFactor` is two, two IPUs are\n",
    "required."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### `randomSeed`\n",
    "\n",
    "An advantage of the IPU architecture is an on-device pseudo-random number\n",
    "generator (PRNG). This option sets both the seed for the PRNG on the IPU\n",
    "and PyTorch's seed, which is usually set using `torch.manual_seed`."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### `useIpuModel`\n",
    "\n",
    "An IPU Model is a simulation, running on a CPU, of an actual IPU. This can be\n",
    "helpful if you're working in an environment where no IPUs are available but\n",
    "still need to make progress on your code. However, the IPU Model doesn't\n",
    "fully support replicated graphs and its numerical results can be slightly\n",
    "different from what you would get with an actual IPU. You can learn more\n",
    "about the IPU Model and its limitations with our\n",
    "[documentation](https://docs.graphcore.ai/projects/poplar-user-guide/en/3.1.0/poplar_programs.html?highlight=ipu%20model#programming-with-poplar)."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## How to set the options\n",
    "\n",
    "These options are callable, and can be chained as they return the instance. One\n",
    "can therefore do the following:"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "opts = (\n",
    "    poptorch.Options()\n",
    "    .deviceIterations(20)\n",
    "    .replicationFactor(2)\n",
    "    .randomSeed(123)\n",
    "    .useIpuModel(True)\n",
    ")"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Summary\n",
    "In this tutorial, we started learning about PopTorch and how to use it to create and train\n",
    "a model which we then be used for inference. We used Torchvision to load the\n",
    "FashionMNIST dataset, a more complex alternative to the well-known MNIST\n",
    "dataset. We created a classification model and trained it on a training subset\n",
    "of our dataset by feeding the images into our model using the PopTorch\n",
    "`DataLoader` class. After training, we evaluated our model using the testing dataset to\n",
    "measure a metric for our model. Lastly, we used the trained model\n",
    "to predict the class of clothing for our own images, which the model had never\n",
    "seen before. When submitting our own images for classification, we needed to make\n",
    "sure that we were aware of the limitations of our model."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Next steps:\n",
    "To continue learning about PopTorch and how to use it to run your models on the\n",
    "IPU, look at our tutorials on [efficient data\n",
    "loading](../efficient_data_loading), [pipelining your model in\n",
    "PopTorch](../pipelining) and [fine-tuning a BERT model to run on the\n",
    "IPU](../finetuning_bert)."
   ],
   "metadata": {}
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "traceability": {
   "sdk_version": "3.0.0+1145",
   "source_file": "walkthrough.py",
   "sst_version": "0.0.8",
   "timestamp": "2022-09-27T15:26"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}