{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5cbba845",
   "metadata": {},
   "source": [
    "Copyright (c) 2020 Graphcore Ltd. All rights reserved."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e1138a2",
   "metadata": {},
   "source": [
    "*Notebook autogenerated from mnist_poptorch.py on 27-Sep-2022*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c98d2c8",
   "metadata": {},
   "source": [
    "# PyTorch (PopTorch) MNIST Training Demo\n",
    "\n",
    "This example demonstrates how to train a neural network for classification on the MNIST dataset using PopTorch.\n",
    "To learn more about PopTorch, see our [PyTorch for the IPU: User Guide](https://docs.graphcore.ai/projects/poptorch-user-guide/en/latest/)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc9ad7f8",
   "metadata": {},
   "source": [
    "## How to use this demo\n",
    "Requirements:\n",
    "- A Poplar SDK environment enabled, with PopTorch installed (see the [Getting Started](https://docs.graphcore.ai/en/latest/getting-started.html) guide for your IPU system)\n",
    "- Python packages installed with `python -m pip install -r requirements.txt`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c5694e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install  -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "330692b7",
   "metadata": {},
   "source": [
    "To run the Jupyter notebook version of this tutorial:\n",
    "1. Enable a Poplar SDK environment and install required packages with `python -m pip install -r requirements.txt`\n",
    "2. In the same environment, install the Jupyter notebook server: `python -m pip install jupyter`\n",
    "3. Launch a Jupyter Server on a specific port: `jupyter-notebook --no-browser --port <port number>`\n",
    "4. Connect via SSH to your remote machine, forwarding your chosen port:\n",
    "`ssh -NL <port number>:localhost:<port number> <your username>@<remote machine>`\n",
    "\n",
    "For more details about this process, or if you need troubleshooting, see our [guide on using IPUs from Jupyter notebooks](../../../tutorials/standard_tools/using_jupyter/README.md)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf6f0753",
   "metadata": {},
   "source": [
    "## Training a PopTorch model for MNIST classification"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d0c8d98",
   "metadata": {},
   "source": [
    "### Importing required libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d2a9772",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "import poptorch\n",
    "import torch.optim as optim"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f4ec513",
   "metadata": {},
   "source": [
    "### Setting hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "711b2009",
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate = 0.03\n",
    "\n",
    "epochs = 10\n",
    "\n",
    "batch_size = 8\n",
    "\n",
    "test_batch_size = 80"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3da823cd",
   "metadata": {},
   "source": [
    "Device iteration defines the number of iterations the device should\n",
    "run over the data before returning to the user.\n",
    "This is equivalent to running the IPU in a loop over that the specified\n",
    "number of iterations, with a new batch of data each time. However, increasing\n",
    "deviceIterations is more efficient because the loop runs on the IPU directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3848cac",
   "metadata": {},
   "outputs": [],
   "source": [
    "device_iterations = 50"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "433ccddd",
   "metadata": {},
   "source": [
    "### Preparing the data\n",
    "We use the `torchvision` package to get the MNIST dataset and we create two data loaders: one for training, one for testing.\n",
    "Source: [The MNIST Database](http://yann.lecun.com/exdb/mnist/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05500b84",
   "metadata": {},
   "outputs": [],
   "source": [
    "local_dataset_path = \"~/.torch/datasets\"\n",
    "\n",
    "transform_mnist = torchvision.transforms.Compose(\n",
    "    [\n",
    "        torchvision.transforms.ToTensor(),\n",
    "        torchvision.transforms.Normalize((0.1307,), (0.3081,)),\n",
    "    ]\n",
    ")\n",
    "\n",
    "training_dataset = torchvision.datasets.MNIST(\n",
    "    local_dataset_path, train=True, download=True, transform=transform_mnist\n",
    ")\n",
    "\n",
    "test_dataset = torchvision.datasets.MNIST(\n",
    "    local_dataset_path, train=False, download=True, transform=transform_mnist\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a160910d",
   "metadata": {},
   "source": [
    "We use the [data loader provided by\n",
    "PopTorch](https://docs.graphcore.ai/projects/poptorch-user-guide/en/latest/pytorch_to_poptorch.html#preparing-your-data).\n",
    "More information about the use of `poptorch.Dataloader` can be found in\n",
    "[PopTorch tutorial on efficient data\n",
    "loading](../../../tutorials/pytorch/efficient_data_loading)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a41bcd7c",
   "metadata": {},
   "source": [
    "A `poptorch.Options()` instance contains a set of default hyperparameters and options for the IPU.\n",
    "This is used by the model and the PopTorch `DataLoader`.\n",
    "To accelerate the training here, we change the default value of\n",
    "[deviceIterations](https://docs.graphcore.ai/projects/poptorch-user-guide/en/latest/batching.html#poptorch-options-deviceiterations)\n",
    "to 50.\n",
    "With that setting the data loader will pick 50 batches of data per step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05ce3b37",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_opts = poptorch.Options()\n",
    "training_opts = training_opts.deviceIterations(device_iterations)\n",
    "\n",
    "training_data = poptorch.DataLoader(\n",
    "    options=training_opts,\n",
    "    dataset=training_dataset,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True,\n",
    "    drop_last=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f86066b",
   "metadata": {},
   "source": [
    "For the validation, we choose not to change `deviceIterations`, we will use a default `poptorch.Options()` instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34a7802d",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data = poptorch.DataLoader(\n",
    "    options=poptorch.Options(),\n",
    "    dataset=test_dataset,\n",
    "    batch_size=test_batch_size,\n",
    "    shuffle=True,\n",
    "    drop_last=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c943cd7",
   "metadata": {},
   "source": [
    "### Defining the model\n",
    "Let's define our neural network.\n",
    "This step is similar to what we would do if we were not using an IPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0618ea47",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Block(nn.Module):\n",
    "    def __init__(self, in_channels, num_filters, kernel_size, pool_size):\n",
    "        super(Block, self).__init__()\n",
    "        self.conv = nn.Conv2d(in_channels, num_filters, kernel_size=kernel_size)\n",
    "        self.pool = nn.MaxPool2d(kernel_size=pool_size)\n",
    "        self.relu = nn.ReLU()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv(x)\n",
    "        x = self.pool(x)\n",
    "        x = self.relu(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class Network(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Network, self).__init__()\n",
    "        self.layer1 = Block(1, 32, 3, 2)\n",
    "        self.layer2 = Block(32, 64, 3, 2)\n",
    "        self.layer3 = nn.Linear(1600, 128)\n",
    "        self.layer3_act = nn.ReLU()\n",
    "        self.layer3_dropout = torch.nn.Dropout(0.5)\n",
    "        self.layer4 = nn.Linear(128, 10)\n",
    "        self.softmax = nn.Softmax(1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.layer1(x)\n",
    "        x = self.layer2(x)\n",
    "        # Flatten layer\n",
    "        x = x.view(-1, 1600)\n",
    "        x = self.layer3_act(self.layer3(x))\n",
    "        x = self.layer4(self.layer3_dropout(x))\n",
    "        x = self.softmax(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03a973a8",
   "metadata": {},
   "source": [
    "To ensure the loss computation is placed on the IPU, we need to set\n",
    "it in the `forward()` method of our `torch.nn.Module`.\n",
    "We define a thin wrapper around the `torch.nn.Module` that will use\n",
    "the cross-entropy loss function.\n",
    "This class is creating a custom module to compose the Neural Network and\n",
    "the Cross Entropy module into one object, which under the hood will invoke\n",
    "the `__call__` function on `nn.Module` and consequently the `forward` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d31abb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TrainingModelWithLoss(torch.nn.Module):\n",
    "    def __init__(self, model):\n",
    "        super().__init__()\n",
    "        self.model = model\n",
    "        self.loss = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "    def forward(self, args, labels=None):\n",
    "        output = self.model(args)\n",
    "        if labels is None:\n",
    "            return output\n",
    "        else:\n",
    "            loss = self.loss(output, labels)\n",
    "            return output, loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61d706c0",
   "metadata": {},
   "source": [
    "Let's initialise the neural network from our defined classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43548220",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Network()\n",
    "model_with_loss = TrainingModelWithLoss(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "687ee903",
   "metadata": {},
   "source": [
    "### From PyTorch to PopTorch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8af2689a",
   "metadata": {},
   "source": [
    "We start by initializing the training `Options` that will be used by the training model.\n",
    "`deviceIterations` is set as we did for the training `DataLoader`.\n",
    "With this setting, the training loop on the device will perform 50 iterations before returning control to the host process.\n",
    "Therefore, from the host point of view, each step will consume 50 batches and perform 50 weight updates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01daff80",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_opts = poptorch.Options().deviceIterations(device_iterations)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cdeafe4",
   "metadata": {},
   "source": [
    "Now let's set the [`OutputMode`](https://docs.graphcore.ai/projects/poptorch-user-guide/en/latest/reference.html#poptorch.OutputMode)\n",
    "for our training. By default, PopTorch will\n",
    "return to the host machine only a limited set of information for performance\n",
    "reasons. This is represented by having `OutputMode.Final` as the default, which\n",
    "means that only the final batch of the internal training loop is returned to the host.\n",
    "When inspecting the training performance as it is executing, values like\n",
    "accuracy or losses will only be returned for that last batch.\n",
    "We can set this to `OutputMode.All` to be able to present the full information.\n",
    "This has an impact on the speed of training, due to overhead of transferring\n",
    "more data to the host machine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "191f28b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_opts = training_opts.outputMode(poptorch.OutputMode.All)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af8010d1",
   "metadata": {},
   "source": [
    "We can check if the model is assembled correctly by printing the string\n",
    "representation of the model object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24048b3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(model_with_loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0521ac8b",
   "metadata": {},
   "source": [
    "Now we apply the model wrapping function, which will perform a shallow copy\n",
    "of the PyTorch model. To train the model, we will use [SGD](https://docs.graphcore.ai/projects/poptorch-user-guide/en/latest/reference.html#poptorch.optim.SGD),\n",
    "the Stochastic Gradient Descent with no momentum.\n",
    "This is also where we pass the `training_opts` defined sooner."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4976137",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_model = poptorch.trainingModel(\n",
    "    model_with_loss,\n",
    "    training_opts,\n",
    "    optimizer=optim.SGD(model.parameters(), lr=learning_rate),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ceb55383",
   "metadata": {},
   "source": [
    "### Training loop\n",
    "We are ready to start training! However to track the accuracy while training\n",
    "we need to define one more helper function. During the training, not every\n",
    "samples prediction is returned for efficiency reasons, so this helper function\n",
    "will check accuracy for labels where prediction is available. This behaviour\n",
    "is controlled by setting `AnchorMode` in `poptorch.Options()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4c8d47e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from metrics import accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6871181a",
   "metadata": {},
   "source": [
    "This code will perform the training over the requested amount of epochs\n",
    "and batches using the configured Graphcore IPUs.\n",
    "Since we set `device_iterations` to 50 and PopTorch `OutputMode` to `All`,\n",
    "`losses` contains the losses of the 50 batches processed by the internal loop.\n",
    "The first call to `training_model` will compile the model for the IPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "988511f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "nr_steps = len(training_data)\n",
    "\n",
    "for epoch in tqdm(range(1, epochs + 1), leave=True, desc=\"Epochs\", total=epochs):\n",
    "    with tqdm(training_data, total=nr_steps, leave=False) as bar:\n",
    "        for data, labels in bar:\n",
    "            preds, losses = training_model(data, labels)\n",
    "\n",
    "            mean_loss = torch.mean(losses).item()\n",
    "\n",
    "            acc = accuracy(preds, labels)\n",
    "            bar.set_description(f\"Loss: {mean_loss:0.4f} | Accuracy: {acc:05.2f}% \")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7dade8a6",
   "metadata": {},
   "source": [
    "We could also do it separately by calling `training_model.compile()` in the first place."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cea1e59b",
   "metadata": {},
   "source": [
    "Now let's release resources so we can reuse them for our validation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee40ec68",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_model.detachFromDevice()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b563e537",
   "metadata": {},
   "source": [
    "## Evaluating the trained model\n",
    "Let's check the validation loss on IPU using the trained model. The weights\n",
    "in `model.parameters()` will be copied from the IPU to the host. The weights\n",
    "from the trained model will be reused to compile the new inference model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8260d6e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "inference_model = poptorch.inferenceModel(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46e98d87",
   "metadata": {},
   "source": [
    "Perform validation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca221226",
   "metadata": {},
   "outputs": [],
   "source": [
    "nr_steps = len(test_data)\n",
    "sum_acc = 0.0\n",
    "with tqdm(test_data, total=nr_steps, leave=False) as bar:\n",
    "    for data, labels in bar:\n",
    "        output = inference_model(data)\n",
    "        sum_acc += accuracy(output, labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40d91b98",
   "metadata": {},
   "source": [
    "Finally the accuracy on the test set is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff2cd4c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Accuracy on test set: {sum_acc / len(test_data):0.2f}%\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}