{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {},
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "##### Copyright 2020 The IREE Authors"
      ],
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Licensed under the Apache License v2.0 with LLVM Exceptions.\n",
        "# See https://llvm.org/LICENSE.txt for license information.\n",
        "# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception"
      ],
      "metadata": {},
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Training and Executing an MNIST Model with IREE"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Overview\n",
        "\n",
        "This notebook covers installing IREE and using it to train a simple neural network on the MNIST dataset."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 1. Install dependencies"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%capture\n",
        "!python -m pip install --upgrade tensorflow"
      ],
      "metadata": {},
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "%%capture\n",
        "!python -m pip install --pre iree-base-compiler iree-base-runtime iree-tools-tf -f https://iree.dev/pip-release-links.html"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "import iree.compiler.tf\n",
        "import iree.runtime\n",
        "import tensorflow as tf\n",
        "\n",
        "# Print version information for future notebook users to reference.\n",
        "!iree-compile --version\n",
        "print(\"TensorFlow version: \", tf.__version__)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 2. Import TensorFlow and Other Dependencies"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "from matplotlib import pyplot as plt\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "\n",
        "tf.random.set_seed(91)\n",
        "np.random.seed(91)\n",
        "\n",
        "plt.rcParams[\"font.family\"] = \"monospace\"\n",
        "plt.rcParams[\"figure.figsize\"] = [8, 4.5]\n",
        "plt.rcParams[\"figure.dpi\"] = 150\n",
        "\n",
        "# Print version information for future notebook users to reference.\n",
        "print(\"TensorFlow version: \", tf.__version__)\n",
        "print(\"Numpy version: \", np.__version__)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 3. Load the MNIST Dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# Keras datasets don't provide metadata.\n",
        "NUM_CLASSES = 10\n",
        "NUM_ROWS, NUM_COLS = 28, 28\n",
        "\n",
        "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()\n",
        "\n",
        "# Reshape into grayscale images:\n",
        "x_train = np.reshape(x_train, (-1, NUM_ROWS, NUM_COLS, 1))\n",
        "x_test = np.reshape(x_test, (-1, NUM_ROWS, NUM_COLS, 1))\n",
        "\n",
        "# Rescale uint8 pixel values into float32 values between 0 and 1:\n",
        "x_train = x_train.astype(np.float32) / 255\n",
        "x_test = x_test.astype(np.float32) / 255\n",
        "\n",
        "# IREE doesn't currently support int8 tensors, so we cast them to int32:\n",
        "y_train = y_train.astype(np.int32)\n",
        "y_test = y_test.astype(np.int32)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "print(\"Sample image from the dataset:\")\n",
        "sample_index = np.random.randint(x_train.shape[0])\n",
        "plt.figure(figsize=(5, 5))\n",
        "plt.imshow(x_train[sample_index].reshape(NUM_ROWS, NUM_COLS), cmap=\"gray\")\n",
        "plt.title(f\"Sample #{sample_index}, label: {y_train[sample_index]}\")\n",
        "plt.axis(\"off\")\n",
        "plt.tight_layout()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 4. Create a Simple DNN"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "MLIR-HLO (the MLIR dialect we use to convert TensorFlow models into assembly that IREE can compile) does not currently support training with a dynamic number of examples, so we compile the model with a fixed batch size (by specifying the batch size in the `tf.TensorSpec`s)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "BATCH_SIZE = 32"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "class TrainableDNN(tf.Module):\n",
        "\n",
        "  def __init__(self):\n",
        "    super().__init__()\n",
        "\n",
        "    # Create a Keras model to train.\n",
        "    inputs = tf.keras.layers.Input((NUM_COLS, NUM_ROWS, 1))\n",
        "    x = tf.keras.layers.Flatten()(inputs)\n",
        "    x = tf.keras.layers.Dense(128)(x)\n",
        "    x = tf.keras.layers.Activation(\"relu\")(x)\n",
        "    x = tf.keras.layers.Dense(10)(x)\n",
        "    outputs = tf.keras.layers.Softmax()(x)\n",
        "    self.model = tf.keras.Model(inputs, outputs)\n",
        "\n",
        "    # Create a loss function and optimizer to use during training.\n",
        "    self.loss = tf.keras.losses.SparseCategoricalCrossentropy()\n",
        "    self.optimizer = tf.keras.optimizers.SGD(learning_rate=1e-2)\n",
        "\n",
        "  @tf.function(input_signature=[\n",
        "      tf.TensorSpec([BATCH_SIZE, NUM_ROWS, NUM_COLS, 1])  # inputs\n",
        "  ])\n",
        "  def predict(self, inputs):\n",
        "    return self.model(inputs, training=False)\n",
        "\n",
        "  # We compile the entire training step by making it a method on the model.\n",
        "  @tf.function(input_signature=[\n",
        "      tf.TensorSpec([BATCH_SIZE, NUM_ROWS, NUM_COLS, 1]),  # inputs\n",
        "      tf.TensorSpec([BATCH_SIZE], tf.int32)  # labels\n",
        "  ])\n",
        "  def learn(self, inputs, labels):\n",
        "    # Capture the gradients from forward prop...\n",
        "    with tf.GradientTape() as tape:\n",
        "      probs = self.model(inputs, training=True)\n",
        "      loss = self.loss(labels, probs)\n",
        "\n",
        "    # ...and use them to update the model's weights.\n",
        "    variables = self.model.trainable_variables\n",
        "    gradients = tape.gradient(loss, variables)\n",
        "    self.optimizer.apply_gradients(zip(gradients, variables))\n",
        "    return loss"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 5. Compile the Model with IREE"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "tf.keras adds a large number of methods to TrainableDNN, and most of them\n",
        "cannot be compiled with IREE. To get around this we tell IREE exactly which\n",
        "methods we would like it to compile."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "exported_names = [\"predict\", \"learn\"]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Choose one of IREE's three backends to compile to. (*Note: Using Vulkan requires installing additional drivers.*)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "backend_choice = \"llvm-cpu (CPU)\" #@param [ \"vmvx (CPU)\", \"llvm-cpu (CPU)\", \"vulkan-spirv (GPU/SwiftShader – requires additional drivers) \" ]\n",
        "backend_choice = backend_choice.split(' ')[0]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# Compile the TrainableDNN module\n",
        "vm_flatbuffer = iree.compiler.tf.compile_module(\n",
        "    TrainableDNN(),\n",
        "    target_backends=[backend_choice],\n",
        "    exported_names=exported_names)\n",
        "compiled_model = iree.runtime.load_vm_flatbuffer(\n",
        "    vm_flatbuffer,\n",
        "    backend=backend_choice)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 6. Train the Compiled Model on MNIST"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "This compiled model is portable, demonstrating that IREE can be used for training on a mobile device. On mobile, IREE has a ~1000 fold binary size advantage over the current TensorFlow solution (which is to use the now-deprecated TF Mobile, as TFLite does not support training at this time)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "#@title Benchmark inference and training\n",
        "print(\"Inference latency:\\n  \", end=\"\")\n",
        "%timeit -n 100 compiled_model.predict(x_train[:BATCH_SIZE])\n",
        "print(\"Training latancy:\\n  \", end=\"\")\n",
        "%timeit -n 100 compiled_model.learn(x_train[:BATCH_SIZE], y_train[:BATCH_SIZE])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# Run the core training loop.\n",
        "losses = []\n",
        "\n",
        "step = 0\n",
        "max_steps = x_train.shape[0] // BATCH_SIZE\n",
        "\n",
        "for batch_start in range(0, x_train.shape[0], BATCH_SIZE):\n",
        "  if batch_start + BATCH_SIZE > x_train.shape[0]:\n",
        "    continue\n",
        "\n",
        "  inputs = x_train[batch_start:batch_start + BATCH_SIZE]\n",
        "  labels = y_train[batch_start:batch_start + BATCH_SIZE]\n",
        "\n",
        "  loss = compiled_model.learn(inputs, labels).to_host()\n",
        "  losses.append(loss)\n",
        "\n",
        "  step += 1\n",
        "  print(f\"\\rStep {step:4d}/{max_steps}: loss = {loss:.4f}\", end=\"\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "#@title Plot the training results\n",
        "!python -m pip install bottleneck\n",
        "import bottleneck as bn\n",
        "smoothed_losses = bn.move_mean(losses, 32)\n",
        "x = np.arange(len(losses))\n",
        "\n",
        "plt.plot(x, smoothed_losses, linewidth=2, label='loss (moving average)')\n",
        "plt.scatter(x, losses, s=16, alpha=0.2, label='loss (per training step)')\n",
        "\n",
        "plt.ylim(0)\n",
        "plt.legend(frameon=True)\n",
        "plt.xlabel(\"training step\")\n",
        "plt.ylabel(\"cross-entropy\")\n",
        "plt.title(\"training loss\");"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 7. Evaluate on Heldout Test Examples"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "#@title Evaluate the network on the test data.\n",
        "accuracies = []\n",
        "\n",
        "step = 0\n",
        "max_steps = x_test.shape[0] // BATCH_SIZE\n",
        "\n",
        "for batch_start in range(0, x_test.shape[0], BATCH_SIZE):\n",
        "  if batch_start + BATCH_SIZE > x_test.shape[0]:\n",
        "    continue\n",
        "\n",
        "  inputs = x_test[batch_start:batch_start + BATCH_SIZE]\n",
        "  labels = y_test[batch_start:batch_start + BATCH_SIZE]\n",
        "\n",
        "  prediction = compiled_model.predict(inputs).to_host()\n",
        "  prediction = np.argmax(prediction, -1)\n",
        "  accuracies.append(np.sum(prediction == labels) / BATCH_SIZE)\n",
        "\n",
        "  step += 1\n",
        "  print(f\"\\rStep {step:4d}/{max_steps}\", end=\"\")\n",
        "print()\n",
        "\n",
        "accuracy = np.mean(accuracies)\n",
        "print(f\"Test accuracy: {accuracy:.3f}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "#@title Display inference predictions on a random selection of heldout data\n",
        "rows = 4\n",
        "columns = 4\n",
        "images_to_display = rows * columns\n",
        "assert BATCH_SIZE >= images_to_display\n",
        "\n",
        "random_index = np.arange(x_test.shape[0])\n",
        "np.random.shuffle(random_index)\n",
        "x_test = x_test[random_index]\n",
        "y_test = y_test[random_index]\n",
        "\n",
        "predictions = compiled_model.predict(x_test[:BATCH_SIZE]).to_host()\n",
        "predictions = np.argmax(predictions, -1)\n",
        "\n",
        "fig, axs = plt.subplots(rows, columns)\n",
        "\n",
        "for i, ax in enumerate(np.ndarray.flatten(axs)):\n",
        "  ax.imshow(x_test[i, :, :, 0])\n",
        "  color = \"#000000\" if predictions[i] == y_test[i] else \"#ff7f0e\"\n",
        "  ax.set_xlabel(f\"prediction={predictions[i]}\", color=color)\n",
        "  ax.grid(False)\n",
        "  ax.set_yticks([])\n",
        "  ax.set_xticks([])\n",
        "\n",
        "fig.tight_layout()"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}