| { |
| "nbformat": 4, |
| "nbformat_minor": 0, |
| "metadata": { |
| "colab": { |
| "name": "mnist_tensorflow.ipynb", |
| "provenance": [], |
| "collapsed_sections": [ |
| "PZtRtMMUZHJS" |
| ] |
| }, |
| "kernelspec": { |
| "name": "python3", |
| "display_name": "Python 3" |
| } |
| }, |
| "cells": [ |
| { |
| "cell_type": "markdown", |
| "metadata": { |
| "id": "PZtRtMMUZHJS", |
| "colab_type": "text" |
| }, |
| "source": [ |
| "##### Copyright 2020 Google LLC.\n", |
| "\n", |
| "Licensed under the Apache License, Version 2.0 (the \"License\");" |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "metadata": { |
| "id": "TouZL3JZZSQe", |
| "colab_type": "code", |
| "colab": {}, |
| "cellView": "both" |
| }, |
| "source": [ |
| "#@title License header\n", |
| "# Copyright 2020 Google LLC\n", |
| "#\n", |
| "# Licensed under the Apache License, Version 2.0 (the \"License\");\n", |
| "# you may not use this file except in compliance with the License.\n", |
| "# You may obtain a copy of the License at\n", |
| "#\n", |
| "# https://www.apache.org/licenses/LICENSE-2.0\n", |
| "#\n", |
| "# Unless required by applicable law or agreed to in writing, software\n", |
| "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", |
| "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", |
| "# See the License for the specific language governing permissions and\n", |
| "# limitations under the License." |
| ], |
| "execution_count": 0, |
| "outputs": [] |
| }, |
| { |
| "cell_type": "markdown", |
| "metadata": { |
| "id": "O6c3qfq5Zv57", |
| "colab_type": "text" |
| }, |
| "source": [ |
| "# MNIST Model TensorFlow Training, IREE Execution\n", |
| "\n", |
| "## Overview\n", |
| "\n", |
| "This notebook creates and trains a TensorFlow 2.0 model for recognizing handwritten digits using the [MNIST dataset](https://en.wikipedia.org/wiki/MNIST_database), then compiles and executes that trained model using IREE.\n", |
| "\n", |
| "## Running Locally\n", |
| "\n", |
| "* Refer to [iree/docs/using_colab.md](https://github.com/google/iree/blob/master/docs/using_colab.md) for general information\n", |
| "* Ensure that you have a recent version of TensorFlow 2.0 [installed on your system](https://www.tensorflow.org/install)\n", |
| "* Enable IREE/TF integration by adding to your user.bazelrc: `build --define=iree_tensorflow=true`\n", |
| "* Start colab by running `python build_tools/scripts/start_colab_kernel.py` (see that file for additional instructions)\n", |
| "* Note: you may need to restart your runtime in order to re-run certain cells. Some of the APIs are not yet stable enough for repeated invocations" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "metadata": { |
| "id": "wBXlE69Ia2QU", |
| "colab_type": "text" |
| }, |
| "source": [ |
| "# Setup Steps" |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "metadata": { |
| "id": "EPF7RGQDYK-M", |
| "colab_type": "code", |
| "colab": {} |
| }, |
| "source": [ |
| "import os\n", |
| "import numpy as np\n", |
| "import tensorflow as tf\n", |
| "from matplotlib import pyplot as plt\n", |
| "from pyiree import compiler as ireec\n", |
| "from pyiree import rt as ireert\n", |
| "\n", |
| "tf.compat.v1.enable_eager_execution()\n", |
| "\n", |
| "SAVE_PATH = os.path.join(os.environ[\"HOME\"], \"saved_models\")\n", |
| "os.makedirs(SAVE_PATH, exist_ok=True)" |
| ], |
| "execution_count": 0, |
| "outputs": [] |
| }, |
| { |
| "cell_type": "code", |
| "metadata": { |
| "id": "43BH_9YcsGs8", |
| "colab_type": "code", |
| "colab": { |
| "base_uri": "https://localhost:8080/", |
| "height": 34 |
| }, |
| "cellView": "form", |
| "outputId": "e24a17ea-297b-491b-f97c-244b411ea8d1" |
| }, |
| "source": [ |
| "#@title Notebook settings { run: \"auto\" }\n", |
| "\n", |
| "#@markdown -----\n", |
| "#@markdown ### Configuration\n", |
| "\n", |
| "backend_choice = \"GPU (vulkan-spirv)\" #@param [ \"GPU (vulkan-spirv)\", \"CPU (interpreter bytecode)\" ]\n", |
| "\n", |
| "if backend_choice == \"GPU (vulkan-spirv)\":\n", |
| " backend_name = \"vulkan-spirv\"\n", |
| " driver_name = \"vulkan\"\n", |
| "else:\n", |
| " backend_name = \"interpreter-bytecode\"\n", |
| " driver_name = \"interpreter\"\n", |
| "tf.print(\"Using IREE compiler backend '%s' and runtime driver '%s'\" % (backend_name, driver_name))\n", |
| "\n", |
| "#@markdown -----\n", |
| "#@markdown ### Training Parameters\n", |
| "\n", |
| "#@markdown <sup>Batch size used to subdivide the training and evaluation samples</sup>\n", |
| "batch_size = 200 #@param { type: \"slider\", min: 10, max: 400 }\n", |
| "\n", |
| "#@markdown <sup>Epochs for training/eval. Higher values take longer to run but generally produce more accurate models</sup>\n", |
| "num_epochs = 5 #@param { type: \"slider\", min: 1, max: 20 }\n", |
| "\n", |
| "#@markdown -----" |
| ], |
| "execution_count": 3, |
| "outputs": [ |
| { |
| "output_type": "stream", |
| "text": [ |
| "Using IREE compiler backend 'vulkan-spirv' and runtime driver 'vulkan'\n" |
| ], |
| "name": "stdout" |
| } |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "metadata": { |
| "id": "5vkQOMOMbXdy", |
| "colab_type": "text" |
| }, |
| "source": [ |
| "# Create and Train MNIST Model in TensorFlow\n", |
| "\n", |
| "The specific details of the training process here aren't critical to the model compilation and execution through IREE." |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "metadata": { |
| "id": "GXZIrReTbTHN", |
| "colab_type": "code", |
| "colab": { |
| "base_uri": "https://localhost:8080/", |
| "height": 486 |
| }, |
| "cellView": "form", |
| "outputId": "d727261b-7457-47ba-97c6-aca339a9d91f" |
| }, |
| "source": [ |
| "#@title Load MNIST dataset, setup training and evaluation\n", |
| "\n", |
| "NUM_CLASSES = 10 # One per digit [0, 1, 2, ..., 9]\n", |
| "IMG_ROWS, IMG_COLS = 28, 28\n", |
| "\n", |
| "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()\n", |
| "tf.print(\"Loaded MNIST dataset!\")\n", |
| "\n", |
| "x_train = x_train.reshape(x_train.shape[0], IMG_ROWS, IMG_COLS, 1)\n", |
| "x_test = x_test.reshape(x_test.shape[0], IMG_ROWS, IMG_COLS, 1)\n", |
| "input_shape = (IMG_ROWS, IMG_COLS, 1)\n", |
| "\n", |
| "# Scale pixel values from [0, 255] integers to [0.0, 1.0] floats.\n", |
| "x_train = x_train.astype(\"float32\") / 255\n", |
| "x_test = x_test.astype(\"float32\") / 255\n", |
| "\n", |
| "steps_per_epoch = int(x_train.shape[0] / batch_size)\n", |
| "steps_per_eval = int(x_test.shape[0] / batch_size)\n", |
| "\n", |
| "# Convert class vectors to binary class matrices.\n", |
| "y_train = tf.keras.utils.to_categorical(y_train, NUM_CLASSES)\n", |
| "y_test = tf.keras.utils.to_categorical(y_test, NUM_CLASSES)\n", |
| "\n", |
| "# Construct batched datasets for training/evaluation.\n", |
| "train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n", |
| "train_dataset = train_dataset.batch(batch_size, drop_remainder=True)\n", |
| "test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))\n", |
| "test_dataset = test_dataset.batch(batch_size, drop_remainder=True)\n", |
| "\n", |
| "# Create a distribution strategy for the dataset (single machine).\n", |
| "strategy = tf.distribute.experimental.CentralStorageStrategy()\n", |
| "train_dist_ds = strategy.experimental_distribute_dataset(train_dataset)\n", |
| "test_dist_ds = strategy.experimental_distribute_dataset(test_dataset)\n", |
| "\n", |
| "tf.print(\"Configured data for training and evaluation!\")\n", |
| "tf.print(\" sample shape: %s\" % str(x_train[0].shape))\n", |
| "tf.print(\" training samples: %s\" % x_train.shape[0])\n", |
| "tf.print(\" test samples: %s\" % x_test.shape[0])\n", |
| "tf.print(\" epochs: %s\" % num_epochs)\n", |
| "tf.print(\" steps/epoch: %s\" % steps_per_epoch)\n", |
| "tf.print(\" steps/eval : %s\" % steps_per_eval)\n", |
| "\n", |
| "tf.print(\"\")\n", |
| "tf.print(\"Sample image from the dataset:\")\n", |
| "SAMPLE_EXAMPLE_INDEX = 1\n", |
| "sample_image = x_test[SAMPLE_EXAMPLE_INDEX]\n", |
| "sample_image_batch = np.expand_dims(sample_image, axis=0)\n", |
| "sample_label = y_test[SAMPLE_EXAMPLE_INDEX]\n", |
| "plt.imshow(sample_image.reshape(IMG_ROWS, IMG_COLS))\n", |
| "plt.show()\n", |
| "tf.print(\"\\nGround truth labels: %s\" % str(sample_label))" |
| ], |
| "execution_count": 4, |
| "outputs": [ |
| { |
| "output_type": "stream", |
| "text": [ |
| "Loaded MNIST dataset!\n", |
| "INFO:tensorflow:ParameterServerStrategy (CentralStorageStrategy if you are using a single machine) with compute_devices = ('/device:CPU:0',), variable_device = '/device:CPU:0'\n", |
| "Configured data for training and evaluation!\n", |
| " sample shape: (28, 28, 1)\n", |
| " training samples: 60000\n", |
| " test samples: 10000\n", |
| " epochs: 5\n", |
| " steps/epoch: 300\n", |
| " steps/eval : 50\n", |
| "\n", |
| "Sample image from the dataset:\n" |
| ], |
| "name": "stdout" |
| }, |
| { |
| "output_type": "display_data", |
| "data": { |
| "text/plain": [ |
| "<Figure size 432x288 with 1 Axes>" |
| ], |
| "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD4CAYAAAAq5pAIAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4xLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy8li6FKAAANzUlEQVR4nO3df6zV9X3H8dcL5IdFVBiMMSRaLMRiF6G9oXV1m8a1s/xRbLK5ks5hY3O7rG5tQtIat6Q2/RGzVN2WNV1oJaWLP+L8UVlqOpHaOFuCXhwFhLZQhyvsChJuB24ZcK/v/XG/NFe93++5nPM9P+T9fCQ355zv+3y/33eOvvie8/2c7/k4IgTg7Dep2w0A6AzCDiRB2IEkCDuQBGEHkjinkzub6mkxXTM6uUsglf/T/+hknPB4tZbCbvs6SX8nabKkb0bEHVXPn64Zeq+vbWWXACpsjc2ltabfxtueLOlrkj4kaamk1baXNrs9AO3Vymf2FZL2RcSLEXFS0gOSVtXTFoC6tRL2BZJ+MebxgWLZ69jutz1ge+CUTrSwOwCtaPvZ+IhYFxF9EdE3RdPavTsAJVoJ+0FJC8c8vqhYBqAHtRL25yQttv1221MlfVTSxnraAlC3pofeImLY9i2S/lWjQ2/rI+KF2joDUKuWxtkj4nFJj9fUC4A24uuyQBKEHUiCsANJEHYgCcIOJEHYgSQIO5AEYQeSIOxAEoQdSIKwA0kQdiAJwg4k0dGfkkZz9n/pysr6yPTyyTnnXv5K5bpbrni4qZ5Ou/T7H6+sz3z23NLavL//UUv7xpnhyA4kQdiBJAg7kARhB5Ig7EAShB1IgrADSTDO3gOGvru4sr5r2T+0bd+nyofoJ+Qn13yzsn5v3/zS2oObfq9y3ZE9e5vqCePjyA4kQdiBJAg7kARhB5Ig7EAShB1IgrADSTDO3gGNxtF/uOyBtu37H3+5qLJ+15YPVNYvubj6evgnlj5SWf/YzMHS2pdvmlO57qLPMc5ep5bCbnu/pOOSRiQNR0RfHU0BqF8dR/ZrIuJIDdsB0EZ8ZgeSaDXsIekJ29ts94/3BNv9tgdsD5zSiRZ3B6BZrb6NvyoiDtr+dUmbbP8kIp4e+4SIWCdpnSSd79ktXnYBoFktHdkj4mBxe1jSo5JW1NEUgPo1HXbbM2zPPH1f0gcl7aqrMQD1auVt/DxJj9o+vZ37IuJ7tXT1FjN87Xsq69+/4msNtjClsvq3Q0sq60/9ccWI538drlx3ydBAZX3S9OmV9a9s/a3K+m1zdpbWhmcNV66LejUd9oh4UdIVNfYCoI0YegOSIOxAEoQdSIKwA0kQdiAJLnGtwasLplbWJzX4N7XR0NoPPlw9vDXy4k8r663Y94XllfX7Zt/ZYAvTSisXfY9jTSfxagNJEHYgCcIOJEHYgSQIO5AEYQeSIOxAEoyz1+DCb2+prP/hwJ9U1j10rLI+PLj/DDuqzydWPllZP29S+Tg6egtHdiAJwg4kQdiBJAg7kARhB5Ig7EAShB1IgnH2DhjZ/bNut1Bq/5evrKzffOFXG2yh+qem1w6+r7Q288k9leuONNgzzgxHdiAJwg4kQdiBJAg7kARhB5Ig7EAShB1IgnH2s9wvb6weR//hn1aPo18wqXocfcuJyZX17V8q/935c489W7ku6tXwyG57ve3DtneNWTbb9ibbe4vbWe1tE0CrJvI2/luSrnvDslslbY6IxZI2F48B9LCGYY+IpyUdfcPiVZI2FPc3SLq+5r4A1KzZz+zzImKwuP+ypHllT7TdL6lfkqbrbU3uDkCrWj4bHxEhKSrq6yKiLyL6plRM8gegvZoN+yHb8yWpuD1cX0sA2qHZsG+UtKa4v0bSY/W0A6BdGn5mt32/pKslzbF9QNLnJd0h6UHbN0t6SdIN7WwSzTvy7tJPWJIaj6M3suYHn6isL/kOY+m9omHYI2J1SenamnsB0EZ8XRZIgrADSRB2IAnCDiRB2IEkuMT1LHBy08WltS2X3dlg7eqhtyu2rKmsv3Ptzyvr/Bx07+DIDiRB2IEkCDuQBGEHkiDsQBKEHUiCsANJMM7+FnDOoksq6198xz+X1mY1uIR124nqfV/8xeqR8pGhoeoNoGdwZAeSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJBhnfwu49MGDlfXlU5v/N3v15j+rrC/58XNNbxu9hSM7kARhB5Ig7EAShB1IgrADSRB2IAnCDiTBOHsPGFpzZWX9C/Ma/fb7tNLKmv2/X7nmOz+7r7LO776fPRoe2W2vt33Y9q4xy263fdD29uJvZXvbBNCqibyN/5ak68ZZfndELCv+Hq+3LQB1axj2iHha0tEO9AKgjVo5QXeL7R3F2/xZZU+y3W97wPbAKTX4wTMAbdNs2L8u6VJJyyQNSio9gxQR6yKiLyL6plScSALQXk2FPSIORcRIRLwm6RuSVtTbFoC6NRV22/PHPPyIpF1lzwXQGxqOs9u+X9LVkubYPiDp85Kutr1MUkjaL+mTbezxLe+cBb9ZWf+dv9xaWT9vUvMff7bsfkdlfckQ16tn0TDsEbF6nMX3tKEXAG3E12WBJAg7kARhB5Ig7EAShB1IgktcO2DPbQsr69/5jX9pafvX7Pyj0hqXsOI0juxAEoQdSIKwA0kQdiAJwg4kQdiBJAg7kATj7B2w7cN3N3hGa7/gc8Gfv1ZaGx4aamnbOHtwZAeSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJBhnPwucmndBaW3KyQUd7OTNRl45UlqLE9XTgXla9fcPJs+d01RPkjQy98LK+t61U5ve9kTEiEtrl/1Fg98gOHasqX1yZAeSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJBhnPwt896H13W6h1G//+3iTAI86cuj8ynVnzT1eWd/6nvua6qnXLf3rWyrriz67pantNjyy215o+ynbu22/YPvTxfLZtjfZ3lvczmqqAwAdMZG38cOS1kbEUknvk/Qp20sl3Sppc0QslrS5eAygRzUMe0QMRsTzxf3jkvZIWiBplaQNxdM2SLq+XU0CaN0ZfWa3fYmk5ZK2SpoXEYNF6WVJ80rW6ZfUL0nT9bZm+wTQogmfjbd9nqSHJX0mIl73TfyICEkx3noRsS4i+iKib0qLP6wIoHkTCrvtKRoN+r0R8Uix+JDt+UV9vqTD7WkRQB0avo23bUn3SNoTEXeNKW2UtEbSHcXtY23p8CywavfHKuub3/VQhzrpvB8tv79r+/7fOFlaOxXlP789ESt33FRZ/+/tzV9+u+CZ4abXrTKRz+zvl3SjpJ22txfLbtNoyB+0fbOklyTd0JYOAdSiYdgj4hlJZVfaX1tvOwDaha/LAkkQdiAJwg4kQdiBJAg7kASXuHbAuX/wH5X1y79SfUljtPG/0szLjlbW23kZ6eX/9vHKevznjJa2v+ihV8uLz+5saduztLelejdwZAeSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJDz6IzOdcb5nx3vNhXJAu2yNzToWR8e9SpUjO5AEYQeSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJAg7kARhB5Ig7EAShB1IgrADSRB2IAnCDiTRMOy2F9p+yvZu2y/Y/nSx/HbbB21vL/5Wtr9dAM2ayPQDw5LWRsTztmdK2mZ7U1G7OyK+2r72ANRlIvOzD0oaLO4ft71H0oJ2NwagXmf0md32JZKWS9paLLrF9g7b623PKlmn3/aA7YFTOtFSswCaN+Gw2z5P0sOSPhMRxyR9XdKlkpZp9Mh/53jrRcS6iOiLiL4pmlZDywCaMaGw256i0aDfGxGPSFJEHIqIkYh4TdI3JK1oX5sAWjWRs/GWdI+kPRFx15jl88c87SOSdtXfHoC6TORs/Psl3Shpp+3txbLbJK22vUxSSNov6ZNt6RBALSZyNv4ZSeP9DvXj9bcDoF34Bh2QBGEHkiDsQBKEHUiCsANJEHYgCcIOJEHYgSQIO5AEYQeSIOxAEoQdSIKwA0kQdiAJR0Tndma/IumlMYvmSDrSsQbOTK/21qt9SfTWrDp7uzgi5o5X6GjY37RzeyAi+rrWQIVe7a1X+5LorVmd6o238UAShB1IotthX9fl/Vfp1d56tS+J3prVkd66+pkdQOd0+8gOoEMIO5BEV8Ju+zrbP7W9z/at3eihjO39tncW01APdLmX9bYP2941Ztls25ts7y1ux51jr0u99cQ03hXTjHf1tev29Ocd/8xue7Kkn0n6gKQDkp6TtDoidne0kRK290vqi4iufwHD9u9KelXStyPiXcWyv5F0NCLuKP6hnBURn+uR3m6X9Gq3p/EuZiuaP3aacUnXS7pJXXztKvq6QR143bpxZF8haV9EvBgRJyU9IGlVF/roeRHxtKSjb1i8StKG4v4Gjf7P0nElvfWEiBiMiOeL+8clnZ5mvKuvXUVfHdGNsC+Q9Isxjw+ot+Z7D0lP2N5mu7/bzYxjXkQMFvdfljSvm82Mo+E03p30hmnGe+a1a2b681Zxgu7NroqId0v6kKRPFW9Xe1KMfgbrpbHTCU3j3SnjTDP+K9187Zqd/rxV3Qj7QUkLxzy+qFjWEyLiYHF7WNKj6r2pqA+dnkG3uD3c5X5+pZem8R5vmnH1wGvXzenPuxH25yQttv1221MlfVTSxi708Sa2ZxQnTmR7hqQPqvemot4oaU1xf42kx7rYy+v0yjTeZdOMq8uvXdenP4+Ijv9JWqnRM/I/l/RX3eihpK9Fkn5c/L3Q7d4k3a/Rt3WnNHpu42ZJvyZps6S9kp6UNLuHevsnSTsl7dBosOZ3qberNPoWfYek7cXfym6/dhV9deR14+uyQBKcoAOSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJP4fcKgKSEIBgPIAAAAASUVORK5CYII=\n" |
| }, |
| "metadata": { |
| "tags": [], |
| "needs_background": "light" |
| } |
| }, |
| { |
| "output_type": "stream", |
| "text": [ |
| "\n", |
| "Ground truth labels: [0. 0. 1. 0. 0. 0. 0. 0. 0. 0.]\n" |
| ], |
| "name": "stdout" |
| } |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "metadata": { |
| "id": "tHq96SIJcNfx", |
| "colab_type": "code", |
| "colab": {}, |
| "cellView": "both" |
| }, |
| "source": [ |
| "#@title Define MNIST model architecture using tf.keras API\n", |
| "\n", |
| "def simple_mnist_model(input_shape):\n", |
| " \"\"\"Creates a simple (multi-layer perceptron) MNIST model.\"\"\"\n", |
| "\n", |
| " model = tf.keras.models.Sequential()\n", |
| " # Flatten to a 1d array (e.g. 28x28 -> 784)\n", |
| " model.add(tf.keras.layers.Flatten(input_shape=input_shape))\n", |
| " # Fully-connected neural layer with 128 neurons, RELU activation\n", |
| " model.add(tf.keras.layers.Dense(128, activation='relu'))\n", |
| " # Fully-connected neural layer returning probability scores for each class\n", |
| " model.add(tf.keras.layers.Dense(10, activation='softmax'))\n", |
| " return model" |
| ], |
| "execution_count": 0, |
| "outputs": [] |
| }, |
| { |
| "cell_type": "code", |
| "metadata": { |
| "id": "7Gdxh7qWcPSO", |
| "colab_type": "code", |
| "colab": { |
| "base_uri": "https://localhost:8080/", |
| "height": 374 |
| }, |
| "cellView": "form", |
| "outputId": "e38d9b12-c27c-4174-824b-8bc232a2c417" |
| }, |
| "source": [ |
| "#@title Train the Keras model\n", |
| "\n", |
| "with strategy.scope():\n", |
| " model = simple_mnist_model(input_shape)\n", |
| " tf.print(\"Constructed Keras MNIST model, training...\")\n", |
| "\n", |
| " optimizer = tf.keras.optimizers.SGD(learning_rate=0.05)\n", |
| " training_loss = tf.keras.metrics.Mean(\"training_loss\", dtype=tf.float32)\n", |
| " training_accuracy = tf.keras.metrics.CategoricalAccuracy(\n", |
| " \"training_accuracy\", dtype=tf.float32)\n", |
| " test_loss = tf.keras.metrics.Mean(\"test_loss\", dtype=tf.float32)\n", |
| " test_accuracy = tf.keras.metrics.CategoricalAccuracy(\n", |
| " \"test_accuracy\", dtype=tf.float32)\n", |
| "\n", |
| " @tf.function\n", |
| " def train_step(iterator):\n", |
| " \"\"\"Training StepFn.\"\"\"\n", |
| "\n", |
| " def step_fn(inputs):\n", |
| " \"\"\"Per-Replica StepFn.\"\"\"\n", |
| " images, labels = inputs\n", |
| " with tf.GradientTape() as tape:\n", |
| " logits = model(images, training=True)\n", |
| " loss = tf.keras.losses.categorical_crossentropy(labels, logits)\n", |
| " loss = tf.reduce_mean(loss) / strategy.num_replicas_in_sync\n", |
| " grads = tape.gradient(loss, model.trainable_variables)\n", |
| " optimizer.apply_gradients(zip(grads, model.trainable_variables))\n", |
| " training_loss.update_state(loss)\n", |
| " training_accuracy.update_state(labels, logits)\n", |
| "\n", |
| " strategy.experimental_run_v2(step_fn, args=(next(iterator),))\n", |
| "\n", |
| " @tf.function\n", |
| " def test_step(iterator):\n", |
| " \"\"\"Evaluation StepFn.\"\"\"\n", |
| "\n", |
| " def step_fn(inputs):\n", |
| " images, labels = inputs\n", |
| " logits = model(images, training=False)\n", |
| " loss = tf.keras.losses.categorical_crossentropy(labels, logits)\n", |
| " loss = tf.reduce_mean(loss) / strategy.num_replicas_in_sync\n", |
| " test_loss.update_state(loss)\n", |
| " test_accuracy.update_state(labels, logits)\n", |
| "\n", |
| " strategy.experimental_run_v2(step_fn, args=(next(iterator),))\n", |
| "\n", |
| " for epoch in range(0, num_epochs):\n", |
| " tf.print(\"Running epoch #%s\" % (epoch + 1))\n", |
| "\n", |
| " train_iterator = iter(train_dist_ds)\n", |
| " for step in range(steps_per_epoch):\n", |
| " train_step(train_iterator)\n", |
| " tf.print(\" Training loss: %f, accuracy: %f\" % (training_loss.result(), training_accuracy.result() * 100))\n", |
| " training_loss.reset_states()\n", |
| " training_accuracy.reset_states()\n", |
| "\n", |
| " test_iterator = iter(test_dist_ds)\n", |
| " for step in range(steps_per_eval):\n", |
| " test_step(test_iterator)\n", |
| " tf.print(\" Test loss : %f, accuracy: %f\" % (test_loss.result(), test_accuracy.result() * 100))\n", |
| " test_loss.reset_states()\n", |
| " test_accuracy.reset_states()\n", |
| "\n", |
| " tf.print(\"Completed training!\")\n", |
| " tf.print(\"\")\n", |
| "\n", |
| " # Run a single prediction on the trained model\n", |
| " tf_prediction = model(sample_image_batch, training=False)\n", |
| " tf.print(\"Sample prediction:\")\n", |
| " tf.print(tf_prediction[0] * 100.0, summarize=100)\n", |
| " tf.print(\"\")" |
| ], |
| "execution_count": 6, |
| "outputs": [ |
| { |
| "output_type": "stream", |
| "text": [ |
| "Constructed Keras MNIST model, training...\n", |
| "Running epoch #1\n", |
| " Training loss: 0.714714, accuracy: 82.386665\n", |
| " Test loss : 0.386100, accuracy: 89.630005\n", |
| "Running epoch #2\n", |
| " Training loss: 0.361406, accuracy: 90.100006\n", |
| " Test loss : 0.312490, accuracy: 91.000000\n", |
| "Running epoch #3\n", |
| " Training loss: 0.308521, accuracy: 91.355003\n", |
| " Test loss : 0.278433, accuracy: 92.000000\n", |
| "Running epoch #4\n", |
| " Training loss: 0.277716, accuracy: 92.261665\n", |
| " Test loss : 0.255600, accuracy: 92.650002\n", |
| "Running epoch #5\n", |
| " Training loss: 0.255256, accuracy: 92.849998\n", |
| " Test loss : 0.237870, accuracy: 93.339996\n", |
| "Completed training!\n", |
| "\n", |
| "Sample prediction:\n", |
| "[0.942700624 0.0189254563 94.8591919 1.60412228 4.52805943e-06 1.21846104 0.845801711 3.33044341e-06 0.510759115 2.41755624e-05]\n", |
| "\n" |
| ], |
| "name": "stdout" |
| } |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "metadata": { |
| "id": "DmespEaFcSEL", |
| "colab_type": "code", |
| "colab": { |
| "base_uri": "https://localhost:8080/", |
| "height": 122 |
| }, |
| "outputId": "07ba7fb8-bd99-495d-941c-17320014573b" |
| }, |
| "source": [ |
| "#@title Export the trained model as a SavedModel, with IREE-compatible settings\n", |
| "\n", |
| "# Since the model was written in sequential style, explicitly wrap in a module.\n", |
| "saved_model_dir = \"/tmp/mnist.sm\"\n", |
| "inference_module = tf.Module()\n", |
| "inference_module.model = model\n", |
| "# Hack: Convert to static shape. Won't be necessary once dynamic shapes are in.\n", |
| "dynamic_input_shape = list(model.inputs[0].shape)\n", |
| "dynamic_input_shape[0] = 1 # Make fixed (batch=1)\n", |
| "# Produce a concrete function.\n", |
| "inference_module.predict = tf.function(\n", |
| " input_signature=[\n", |
| " tf.TensorSpec(dynamic_input_shape, model.inputs[0].dtype)])(\n", |
| " lambda x: model.call(x, training=False))\n", |
| "save_options = tf.saved_model.SaveOptions(save_debug_info=True)\n", |
| "tf.print(\"Exporting SavedModel to %s\" % saved_model_dir)\n", |
| "tf.saved_model.save(inference_module, saved_model_dir, options=save_options)" |
| ], |
| "execution_count": 7, |
| "outputs": [ |
| { |
| "output_type": "stream", |
| "text": [ |
| "Exporting SavedModel to /tmp/mnist.sm\n", |
| "WARNING:tensorflow:From /usr/local/lib/python3.7/dist-packages/tensorflow_core/python/ops/resource_variable_ops.py:1788: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n", |
| "Instructions for updating:\n", |
| "If using Keras pass *_constraint arguments to layers.\n", |
| "INFO:tensorflow:Assets written to: /tmp/mnist.sm/assets\n" |
| ], |
| "name": "stdout" |
| } |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "metadata": { |
| "id": "nZdVUd_dgTtc", |
| "colab_type": "text" |
| }, |
| "source": [ |
| "# Compile and Execute MNIST Model using IREE" |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "metadata": { |
| "id": "rqwIx4j4gS1a", |
| "colab_type": "code", |
| "colab": { |
| "base_uri": "https://localhost:8080/", |
| "height": 717 |
| }, |
| "outputId": "2a9b7627-8836-427d-af51-5e001c6a9226" |
| }, |
| "source": [ |
| "#@title Load the SavedModel into IREE's compiler as MLIR xla_hlo\n", |
| "\n", |
| "compiler_module = ireec.tf_load_saved_model(\n", |
| " saved_model_dir, exported_names=[\"predict\"])\n", |
| "tf.print(\"Imported MLIR:\\n\", compiler_module.to_asm(large_element_limit=100))\n", |
| "\n", |
| "# Write to a file for use outside of this notebook.\n", |
| "mnist_mlir_path = os.path.join(SAVE_PATH, \"mnist.mlir\")\n", |
| "with open(mnist_mlir_path, \"wt\") as output_file:\n", |
| " output_file.write(compiler_module.to_asm())\n", |
| "print(\"Wrote MLIR to path '%s'\" % mnist_mlir_path)" |
| ], |
| "execution_count": 8, |
| "outputs": [ |
| { |
| "output_type": "stream", |
| "text": [ |
| "Imported MLIR:\n", |
| " \n", |
| "\n", |
| "module {\n", |
| " flow.variable @\"__iree_flow___sm_node15__model.layer-2.kernel\" opaque<\"\", \"0xDEADBEEF\"> : tensor<784x128xf32>\n", |
| " flow.variable @\"__iree_flow___sm_node16__model.layer-2.bias\" opaque<\"\", \"0xDEADBEEF\"> : tensor<128xf32>\n", |
| " flow.variable @\"__iree_flow___sm_node21__model.layer-3.kernel\" opaque<\"\", \"0xDEADBEEF\"> : tensor<128x10xf32>\n", |
| " flow.variable @\"__iree_flow___sm_node22__model.layer-3.bias\" dense<[-0.124819458, 0.10169369, -0.0273698159, -0.0204691291, 0.0623821244, 0.213322133, -0.036071334, 0.112993836, -0.235609874, -0.0460525416]> : tensor<10xf32>\n", |
| " func @predict(%arg0: tensor<1x28x28x1xf32>) -> tensor<1x10xf32> attributes {iree.module.export, iree.reflection = {abi = \"sip\", abiv = 1 : i32, sip = \"I8!S5!k0_0R3!_0\"}, tf._input_shapes = [\"tfshape$dim { size: 1 } dim { size: 28 } dim { size: 28 } dim { size: 1 }\", \"tfshape$unknown_rank: true\", \"tfshape$unknown_rank: true\", \"tfshape$unknown_rank: true\", \"tfshape$unknown_rank: true\"], tf.signature.is_stateful} {\n", |
| " %0 = flow.variable.load @\"__iree_flow___sm_node22__model.layer-3.bias\" : tensor<10xf32>\n", |
| " %1 = flow.variable.load @\"__iree_flow___sm_node21__model.layer-3.kernel\" : tensor<128x10xf32>\n", |
| " %2 = flow.variable.load @\"__iree_flow___sm_node16__model.layer-2.bias\" : tensor<128xf32>\n", |
| " %3 = flow.variable.load @\"__iree_flow___sm_node15__model.layer-2.kernel\" : tensor<784x128xf32>\n", |
| " %4 = xla_hlo.constant opaque<\"\", \"0xDEADBEEF\"> : tensor<1x128xf32>\n", |
| " %5 = xla_hlo.constant dense<0xFF800000> : tensor<f32>\n", |
| " %6 = xla_hlo.constant dense<0.000000e+00> : tensor<f32>\n", |
| " %7 = \"xla_hlo.reshape\"(%arg0) : (tensor<1x28x28x1xf32>) -> tensor<1x784xf32>\n", |
| " %8 = \"xla_hlo.dot\"(%7, %3) : (tensor<1x784xf32>, tensor<784x128xf32>) -> tensor<1x128xf32>\n", |
| " %9 = \"xla_hlo.add\"(%8, %2) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<1x128xf32>, tensor<128xf32>) -> tensor<1x128xf32>\n", |
| " %10 = xla_hlo.max %9, %4 : tensor<1x128xf32>\n", |
| " %11 = \"xla_hlo.dot\"(%10, %1) : (tensor<1x128xf32>, tensor<128x10xf32>) -> tensor<1x10xf32>\n", |
| " %12 = \"xla_hlo.add\"(%11, %0) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<1x10xf32>, tensor<10xf32>) -> tensor<1x10xf32>\n", |
| " %13 = \"xla_hlo.reduce\"(%12, %5) ( {\n", |
| " ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):\t// no predecessors\n", |
| " %18 = xla_hlo.max %arg1, %arg2 : tensor<f32>\n", |
| " \"xla_hlo.return\"(%18) : (tensor<f32>) -> ()\n", |
| " }) {dimensions = dense<1> : tensor<1xi64>} : (tensor<1x10xf32>, tensor<f32>) -> tensor<1xf32>\n", |
| " %14 = \"xla_hlo.sub\"(%12, %13) {broadcast_dimensions = dense<0> : tensor<1xi64>} : (tensor<1x10xf32>, tensor<1xf32>) -> tensor<1x10xf32>\n", |
| " %15 = \"xla_hlo.exp\"(%14) : (tensor<1x10xf32>) -> tensor<1x10xf32>\n", |
| " %16 = \"xla_hlo.reduce\"(%15, %6) ( {\n", |
| " ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):\t// no predecessors\n", |
| " %18 = xla_hlo.add %arg1, %arg2 : tensor<f32>\n", |
| " \"xla_hlo.return\"(%18) : (tensor<f32>) -> ()\n", |
| " }) {dimensions = dense<1> : tensor<1xi64>} : (tensor<1x10xf32>, tensor<f32>) -> tensor<1xf32>\n", |
| " %17 = \"xla_hlo.div\"(%15, %16) {broadcast_dimensions = dense<0> : tensor<1xi64>} : (tensor<1x10xf32>, tensor<1xf32>) -> tensor<1x10xf32>\n", |
| " return %17 : tensor<1x10xf32>\n", |
| " }\n", |
| "}\n", |
| "\n", |
| "Wrote MLIR to path '/usr/local/google/home/scotttodd/saved_models/mnist.mlir'\n" |
| ], |
| "name": "stdout" |
| } |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "metadata": { |
| "id": "IDHI7h3khJr9", |
| "colab_type": "code", |
| "colab": { |
| "base_uri": "https://localhost:8080/", |
| "height": 51 |
| }, |
| "outputId": "f0676fc3-d7e7-497e-f28d-480e33c37fc9" |
| }, |
| "source": [ |
| "#@title Compile the xla_hlo MLIR and prepare a context to execute it\n", |
| "\n", |
| "# Compile the MLIR module into a VM module for execution\n", |
| "flatbuffer_blob = compiler_module.compile(target_backends=[backend_name])\n", |
| "vm_module = ireert.VmModule.from_flatbuffer(flatbuffer_blob)\n", |
| "\n", |
| "# Register the module with a runtime context\n", |
| "config = ireert.Config(driver_name)\n", |
| "ctx = ireert.SystemContext(config=config)\n", |
| "ctx.add_module(vm_module)" |
| ], |
| "execution_count": 9, |
| "outputs": [ |
| { |
| "output_type": "stream", |
| "text": [ |
| "Created IREE driver vulkan: <pyiree.rt.binding.HalDriver object at 0x7f6e4471bc70>\n", |
| "SystemContext driver=<pyiree.rt.binding.HalDriver object at 0x7f6e4471bc70>\n" |
| ], |
| "name": "stderr" |
| } |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "metadata": { |
| "id": "SKflpnLtkLYE", |
| "colab_type": "code", |
| "colab": { |
| "base_uri": "https://localhost:8080/", |
| "height": 102 |
| }, |
| "outputId": "48996e9a-f70f-4f45-b53c-fd9fab80d78d" |
| }, |
| "source": [ |
| "#@title Execute the compiled module and compare the results with TensorFlow\n", |
| "\n", |
| "# Invoke the 'predict' function with a single image as an argument\n", |
| "iree_prediction = ctx.modules.module.predict(sample_image_batch)\n", |
| "\n", |
| "tf.print(\"IREE prediction ('%s' backend, '%s' driver):\" % (backend_name, driver_name))\n", |
| "tf.print(tf.convert_to_tensor(iree_prediction[0]) * 100.0, summarize=100)\n", |
| "tf.print(\"\")\n", |
| "tf.print(\"TensorFlow prediction:\")\n", |
| "tf.print(tf_prediction[0] * 100.0, summarize=100)" |
| ], |
| "execution_count": 10, |
| "outputs": [ |
| { |
| "output_type": "stream", |
| "text": [ |
| "IREE prediction ('vulkan-spirv' backend, 'vulkan' driver):\n", |
| "[0.942700088 0.018925447 94.8592 1.60412395 4.52806762e-06 1.2184602 0.845801711 3.33043749e-06 0.510759652 2.41755624e-05]\n", |
| "\n", |
| "TensorFlow prediction:\n", |
| "[0.942700624 0.0189254563 94.8591919 1.60412228 4.52805943e-06 1.21846104 0.845801711 3.33044341e-06 0.510759115 2.41755624e-05]\n" |
| ], |
| "name": "stdout" |
| } |
| ] |
| } |
| ] |
| } |
