{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0e7edd79",
   "metadata": {},
   "source": [
    "Copyright (c) 2020 Graphcore Ltd. All rights reserved."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "78cdcfe1",
   "metadata": {},
   "source": [
    "*Notebook autogenerated from mnist.py on 28-Sep-2022*"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c3411906",
   "metadata": {},
   "source": [
    "# Training a simple TensorFlow 2 model on MNIST with an IPU\n",
    "\n",
    "This tutorial shows how to train a simple model using the MNIST numerical\n",
    "dataset on a single IPU. The dataset consists of 60,000 images of handwritten\n",
    "digits (0-9) that must be classified according to which digit they represent."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a0e229fa",
   "metadata": {},
   "source": [
    "We will do the following steps in order:\n",
    "\n",
    "1. Load and pre-process the MNIST dataset from Keras.\n",
    "2. Define a simple model.\n",
    "3. Configure the IPU system\n",
    "4. Train the model on the IPU."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d0ae782f",
   "metadata": {},
   "source": [
    "## Environment setup\n",
    "\n",
    "The best way to run this demo is on Paperspace Gradient’s cloud IPUs because everything is already set up for you. 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/graphcore/Gradient-Tensorflow2/issues).\n",
    "\n",
    "\n",
    "To run the demo using other IPU hardware, you need to have the Poplar SDK enabled. Refer to the [Getting Started guide](https://docs.graphcore.ai/en/latest/getting-started.html#getting-started) for your system for details on how to enable the Poplar SDK. Also refer to the [Jupyter Quick Start guide](https://docs.graphcore.ai/projects/jupyter-notebook-quick-start/en/latest/index.html) for how to set up Jupyter to be able to run this notebook on a remote IPU machine."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5e744433",
   "metadata": {},
   "source": [
    "## 1. Import the necessary libraries\n",
    "\n",
    "First of all, we need to import APIs that will be used in the example."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "bede273d",
   "metadata": {},
   "source": [
    "\n",
    "In order to improve usability and support for future users, Graphcore would like to collect information about the\n",
    "applications and code being run in this notebook. The following information will be anonymised before being sent to Graphcore:\n",
    "\n",
    "- User progression through the notebook\n",
    "- Notebook details: number of cells, code being run and the output of the cells\n",
    "- Environment details\n",
    "\n",
    "You can disable logging at any time by running `%unload_ext graphcore_cloud_tools.notebook_logging.gc_logger` from any cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67f3092a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -r requirements.txt\n",
    "%load_ext graphcore_cloud_tools.notebook_logging.gc_logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76d1d03c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "from tensorflow import keras\n",
    "from tensorflow.python import ipu\n",
    "\n",
    "if tf.__version__[0] != \"2\":\n",
    "    raise ImportError(\"TensorFlow 2 is required for this example\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "37b43a26",
   "metadata": {},
   "source": [
    "For the `ipu` module to function properly, we must import it directly rather\n",
    "than accessing it through the top-level TensorFlow module.\n",
    "\n",
    "## 2. Prepare the dataset\n",
    "\n",
    "We can access the MNIST dataset through keras:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c04cb961",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_dataset():\n",
    "    mnist = keras.datasets.mnist\n",
    "\n",
    "    (x_train, y_train), _ = mnist.load_data()\n",
    "    x_train = x_train / 255.0\n",
    "\n",
    "    train_ds = (\n",
    "        tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
    "        .shuffle(10000)\n",
    "        .batch(32, drop_remainder=True)\n",
    "    )\n",
    "    train_ds = train_ds.map(\n",
    "        lambda d, l: (tf.cast(d, tf.float32), tf.cast(l, tf.float32))\n",
    "    )\n",
    "\n",
    "    # Create a looped version of the dataset\n",
    "    return train_ds.repeat()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "23bc5340",
   "metadata": {},
   "source": [
    "We normalise the dataset by dividing each element of `x_train` (pixel values)\n",
    "by 255. This results in smaller numbers that range from 0 to 1, which leads\n",
    "to faster computation.\n",
    "\n",
    "Some extra care must be taken when preparing a dataset for training a Keras\n",
    "model on the IPU. The Poplar software stack does not support using tensors\n",
    "with shapes which are not known when the model is compiled.\n",
    "\n",
    "To address this we use the `.batch()` method to make sure the sizes of our\n",
    "dataset are divisible by the batch size. The `.batch()` method takes the batch\n",
    "size as an argument and has the option to discard the remaining elements after\n",
    "the dataset is divided (`drop_remainder`). This option must be\n",
    "set to true in order to use the dataset with Keras model on the IPU.\n",
    "\n",
    "## 3. Define the model\n",
    "\n",
    "Next, we define our model using the Keras Sequential API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e8395c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_model():\n",
    "    model = keras.Sequential(\n",
    "        [\n",
    "            keras.layers.Flatten(),\n",
    "            keras.layers.Dense(128, activation=\"relu\"),\n",
    "            keras.layers.Dense(10, activation=\"softmax\"),\n",
    "        ]\n",
    "    )\n",
    "    return model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "48c4b1b0",
   "metadata": {},
   "source": [
    "## 4. Add IPU configuration\n",
    "\n",
    "To use the IPU, we must create an IPU configuration.\n",
    "We can use `cfg.auto_select_ipus = 1` to automatically select one IPU:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b439e7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure the IPU system\n",
    "cfg = ipu.config.IPUConfig()\n",
    "cfg.auto_select_ipus = 1\n",
    "cfg.configure_ipu_system()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f0835084",
   "metadata": {},
   "source": [
    "This is all we need to get a small model up and running, though a full list of\n",
    "configuration options is available in the [API\n",
    "documentation](https://docs.graphcore.ai/projects/tensorflow-user-guide/en/latest/tensorflow/api.html#tensorflow.python.ipu.config.IPUConfig).\n",
    "\n",
    "If you're interested in learning how to optimally use models that require\n",
    "multiple IPUs (for example due to their size), see the section on pipelining\n",
    "from our documentation on [model\n",
    "parallelism](https://docs.graphcore.ai/projects/tf-model-parallelism/en/latest/model.html).\n",
    "\n",
    "> To see how this process can be implemented, head over to the pipelining\n",
    "section of our [TensorFlow 2 Keras\n",
    "tutorial](../../../tutorials/tensorflow2/keras).\n",
    "\n",
    "## 5. Specify IPU strategy\n",
    "\n",
    "Next, add the following code after the configuration:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99bc3fca",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an IPU distribution strategy.\n",
    "strategy = ipu.ipu_strategy.IPUStrategy()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "da2fb6d5",
   "metadata": {},
   "source": [
    "The `tf.distribute.Strategy` is an API to distribute training and inference\n",
    "across multiple devices. `IPUStrategy` is a subclass which targets a system\n",
    "with one or more IPUs attached. For a multi-system configuration, the\n",
    "[PopDistStrategy](https://docs.graphcore.ai/projects/tensorflow-user-guide/en/latest/tensorflow/api.html#tensorflow.python.ipu.horovod.popdist_strategy.PopDistStrategy)\n",
    "should be used, in conjunction with our PopDist library.\n",
    "\n",
    "> To see an example of how to distribute training and inference over multiple\n",
    "instances with PopDist, head over to our [TensorFlow 2 PopDist\n",
    "example](../../../feature_examples/tensorflow2/popdist).\n",
    "\n",
    "## 6. Wrap the model within the IPU strategy scope\n",
    "\n",
    "Creating variables within the scope of the `IPUStrategy` will ensure that they\n",
    "are placed on the IPU, but the initialization for the variables will be\n",
    "performed on the CPU device. To do this, we create a `strategy.scope()` context\n",
    "manager and put all the model code inside of it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d238987",
   "metadata": {},
   "outputs": [],
   "source": [
    "with strategy.scope():\n",
    "    # Create an instance of the model.\n",
    "    model = create_model()\n",
    "\n",
    "    # Get the training dataset.\n",
    "    ds = create_dataset()\n",
    "\n",
    "    # Train the model.\n",
    "    model.compile(\n",
    "        loss=keras.losses.SparseCategoricalCrossentropy(),\n",
    "        optimizer=keras.optimizers.SGD(),\n",
    "        steps_per_execution=100,\n",
    "    )\n",
    "    model.fit(ds, steps_per_epoch=2000, epochs=4)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7cdc27ad",
   "metadata": {},
   "source": [
    "The `steps_per_execution` argument in `model.compile()` sets the number of\n",
    "batches processed in each execution of the underlying IPU program. Not\n",
    "specifying this argument causes the program that runs on the IPU to only process\n",
    "a single batch per execution, which means more time is wasted waiting for I/O\n",
    "instead of using the IPU.\n",
    "\n",
    "Another way to speed up the training of a model is through replication. This\n",
    "process involves copying the model on each of multiple IPUs, updating the\n",
    "parameters of the model on all IPUs after each forward and backward pass. To\n",
    "learn more about this process, head over to our documentation on [graph\n",
    "replication](https://docs.graphcore.ai/projects/memory-performance-optimisation/en/latest/optimising-performance.html#graph-replication).\n",
    "\n",
    "> To see how this process can be implemented, take a look at the Replication\n",
    "section of our [TensorFlow 2 Keras\n",
    "tutorial](../../../tutorials/tensorflow2/keras).\n",
    "\n",
    "## Other useful resources\n",
    "\n",
    "- [TensorFlow\n",
    "  Docs](https://docs.graphcore.ai/en/latest/software.html#tensorflow): all\n",
    "  Graphcore documentation specifically relating to TensorFlow.\n",
    "\n",
    "- [IPU TensorFlow 2 Code\n",
    "  Examples](../../../feature_examples/tensorflow2):\n",
    "  examples of different use cases of TensorFlow 2 on the IPU."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}