From c674a6ac399c105e549d4a6ed9fb99a3ebe75e6a Mon Sep 17 00:00:00 2001 From: Scott Todd Date: Thu, 31 Aug 2023 15:56:18 -0700 Subject: [PATCH] Update sample Colab notebooks to use tf-nightly. Also adapt to API changes (MLIR bytecode default). --- samples/colab/edge_detection.ipynb | 115 +++++----- samples/colab/mnist_training.ipynb | 194 ++++++++++------- samples/colab/resnet.ipynb | 92 +++++--- samples/colab/test_notebooks.py | 11 +- .../colab/tflite_text_classification.ipynb | 57 +++-- samples/dynamic_shapes/dynamic_shapes.ipynb | 202 ++++++++++-------- .../variables_and_state.ipynb | 184 +++++++++------- 7 files changed, 504 insertions(+), 351 deletions(-) diff --git a/samples/colab/edge_detection.ipynb b/samples/colab/edge_detection.ipynb index e7f74860119a..17f8e983f112 100644 --- a/samples/colab/edge_detection.ipynb +++ b/samples/colab/edge_detection.ipynb @@ -56,6 +56,18 @@ "## Setup" ] }, + { + "cell_type": "code", + "source": [ + "%%capture\n", + "!python -m pip install --upgrade tf-nightly # Needed for stablehlo export in TF>=2.14" + ], + "metadata": { + "id": "Nt3QaBP0s65F" + }, + "execution_count": 2, + "outputs": [] + }, { "cell_type": "code", "metadata": { @@ -63,31 +75,23 @@ "colab": { "base_uri": "https://localhost:8080/" }, - "outputId": "f9054d95-2a2d-479f-f787-178fd0291026" + "outputId": "53aee7ec-e390-4293-f51f-2638b75148d3" }, "source": [ "!python -m pip install iree-compiler iree-runtime iree-tools-tf -f https://openxla.github.io/iree/pip-release-links.html" ], - "execution_count": 2, + "execution_count": 3, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ - "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n", "Looking in links: https://openxla.github.io/iree/pip-release-links.html\n", - "Collecting iree-compiler\n", - " Downloading https://github.com/openxla/iree/releases/download/candidate-20230426.501/iree_compiler-20230426.501-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (54.2 MB)\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m54.2/54.2 MB\u001b[0m \u001b[31m9.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[?25hCollecting iree-runtime\n", - " Downloading https://github.com/openxla/iree/releases/download/candidate-20230426.501/iree_runtime-20230426.501-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (2.6 MB)\n", - "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.6/2.6 MB\u001b[0m \u001b[31m30.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", - "\u001b[?25hCollecting iree-tools-tf\n", - " Downloading https://github.com/openxla/iree/releases/download/candidate-20230426.501/iree_tools_tf-20230426.501-py3-none-linux_x86_64.whl (32 kB)\n", - "Requirement already satisfied: numpy in /usr/local/lib/python3.9/dist-packages (from iree-compiler) (1.22.4)\n", - "Requirement already satisfied: PyYAML in /usr/local/lib/python3.9/dist-packages (from iree-compiler) (6.0)\n", - "Installing collected packages: iree-tools-tf, iree-runtime, iree-compiler\n", - "Successfully installed iree-compiler-20230426.501 iree-runtime-20230426.501 iree-tools-tf-20230426.501\n" + "Requirement already satisfied: iree-compiler in /usr/local/lib/python3.10/dist-packages (20230831.630)\n", + "Requirement already satisfied: iree-runtime in /usr/local/lib/python3.10/dist-packages (20230831.630)\n", + "Requirement already satisfied: iree-tools-tf in /usr/local/lib/python3.10/dist-packages (20230831.630)\n", + "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from iree-compiler) (1.23.5)\n", + "Requirement already satisfied: PyYAML in /usr/local/lib/python3.10/dist-packages (from iree-compiler) (6.0.1)\n" ] } ] @@ -99,7 +103,7 @@ "colab": { "base_uri": "https://localhost:8080/" }, - "outputId": "d0d16ea9-7a0c-45cd-eb6f-36eb6baf170e" + "outputId": "8253ff7c-a1b8-44c6-f17c-b9d93dfe30cd" }, "source": [ "#@title Imports\n", @@ -119,13 +123,13 @@ "# Print version information for future notebook users to reference.\n", "print(\"TensorFlow version: \", tf.__version__)" ], - "execution_count": 15, + "execution_count": 4, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ - "TensorFlow version: 2.12.0\n" + "TensorFlow version: 2.15.0-dev20230831\n" ] } ] @@ -142,7 +146,7 @@ "ARTIFACTS_DIR = os.path.join(tempfile.gettempdir(), \"iree\", \"colab_artifacts\")\n", "os.makedirs(ARTIFACTS_DIR, exist_ok=True)" ], - "execution_count": 4, + "execution_count": 5, "outputs": [] }, { @@ -160,7 +164,7 @@ " sobel_x = tf.constant([[-1.0, 0.0, 1.0],\n", " [-2.0, 0.0, 2.0],\n", " [-1.0, 0.0, 1.0]],\n", - " dtype=tf.float32, shape=[3, 3, 1, 1]) \n", + " dtype=tf.float32, shape=[3, 3, 1, 1])\n", " sobel_y = tf.constant([[ 1.0, 2.0, 1.0],\n", " [ 0.0, 0.0, 0.0],\n", " [-1.0, -2.0, -1.0]],\n", @@ -171,7 +175,7 @@ "\n", "tf_module = EdgeDetectionModule()" ], - "execution_count": 5, + "execution_count": 6, "outputs": [] }, { @@ -180,10 +184,10 @@ "cellView": "form", "colab": { "base_uri": "https://localhost:8080/", - "height": 391 + "height": 356 }, "id": "Q2no8DO_s125", - "outputId": "4eb80633-b2dc-4262-83cf-7794ba9a9c09" + "outputId": "2a427588-3a06-4f6b-b15d-76242b9643ae" }, "source": [ "#@title Load a test image of a [labrador](https://commons.wikimedia.org/wiki/File:YellowLabradorLooking_new.jpg) and run the module with TF\n", @@ -219,16 +223,8 @@ "# Plot the input and output images\n", "show_images(image, tf_edges)" ], - "execution_count": 6, + "execution_count": 7, "outputs": [ - { - "output_type": "stream", - "name": "stdout", - "text": [ - "Downloading data from https://storage.googleapis.com/download.tensorflow.org/example_images/YellowLabradorLooking_new.jpg\n", - "83281/83281 [==============================] - 0s 0us/step\n" - ] - }, { "output_type": "display_data", "data": { @@ -262,7 +258,7 @@ "backend_choice = backend_choice.split(\" \")[0]\n", "backend = module_utils.BackendInfo(backend_choice)" ], - "execution_count": 7, + "execution_count": 8, "outputs": [] }, { @@ -273,7 +269,7 @@ "height": 356 }, "id": "2OPBr0eGroE0", - "outputId": "85c453c7-81ba-4181-f872-c2e3cb8d3b9e" + "outputId": "128e76b9-9dd7-49b7-ee4f-0612c8e25bfd" }, "source": [ "#@title Compile and Run the EdgeDetectionModule with IREE.\n", @@ -283,7 +279,7 @@ "iree_edges = module.edge_detect_sobel_operator(image)\n", "show_images(image, iree_edges)" ], - "execution_count": 8, + "execution_count": 9, "outputs": [ { "output_type": "display_data", @@ -320,13 +316,13 @@ "base_uri": "https://localhost:8080/" }, "id": "6YGqN2uqP_7P", - "outputId": "bb1c4be0-ca97-4fd1-8344-0f6d740a415d" + "outputId": "fe3f580c-bbac-48ec-bbbb-1029fa8edce3" }, "source": [ "#@title Construct a module containing the edge detection function\n", "\n", "# Do *not* further compile to a bytecode module for a particular backend.\n", - "# \n", + "#\n", "# By stopping at stablehlo, we can more easily take advantage of future\n", "# compiler improvements within IREE. For a production application, we would\n", "# probably want to freeze the version of IREE used and compile as completely as\n", @@ -334,12 +330,18 @@ "# the application at runtime.\n", "compiler_module = tfc.compile_module(\n", " EdgeDetectionModule(), import_only=True)\n", - "print(\"Edge Detection MLIR: \", compiler_module)\n", "\n", - "edge_detection_mlir_path = os.path.join(ARTIFACTS_DIR, \"edge_detection.mlir\")\n", - "with open(edge_detection_mlir_path, \"wt\") as output_file:\n", + "# Save the imported MLIR to disk.\n", + "imported_mlirbc_path = os.path.join(ARTIFACTS_DIR, \"edge_detection.mlirbc\")\n", + "with open(imported_mlirbc_path, \"wb\") as output_file:\n", " output_file.write(compiler_module)\n", - "print(f\"Wrote MLIR to path '{edge_detection_mlir_path}'\")" + "print(f\"Wrote MLIR to path '{imported_mlirbc_path}'\")\n", + "\n", + "# Copy MLIR bytecode to MLIR text and see how the compiler views this program.\n", + "imported_mlir_path = os.path.join(ARTIFACTS_DIR, \"edge_detection.mlir\")\n", + "!iree-ir-tool copy {imported_mlirbc_path} -o {imported_mlir_path}\n", + "print(\"Edge Detection MLIR:\")\n", + "!cat {imported_mlir_path}" ], "execution_count": 11, "outputs": [ @@ -347,21 +349,21 @@ "output_type": "stream", "name": "stdout", "text": [ - "Edge Detection MLIR: module {\n", + "Wrote MLIR to path '/tmp/iree/colab_artifacts/edge_detection.mlirbc'\n", + "Edge Detection MLIR:\n", + "module {\n", " func.func @edge_detect_sobel_operator(%arg0: tensor<1x128x128x1xf32>) -> tensor<1x128x128x1xf32> {\n", " %0 = stablehlo.constant dense<[[[[-1.000000e+00]], [[0.000000e+00]], [[1.000000e+00]]], [[[-2.000000e+00]], [[0.000000e+00]], [[2.000000e+00]]], [[[-1.000000e+00]], [[0.000000e+00]], [[1.000000e+00]]]]> : tensor<3x3x1x1xf32>\n", " %1 = stablehlo.constant dense<[[[[1.000000e+00]], [[2.000000e+00]], [[1.000000e+00]]], [[[0.000000e+00]], [[0.000000e+00]], [[0.000000e+00]]], [[[-1.000000e+00]], [[-2.000000e+00]], [[-1.000000e+00]]]]> : tensor<3x3x1x1xf32>\n", - " %2 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x128x128x1xf32>, tensor<3x3x1x1xf32>) -> tensor<1x128x128x1xf32>\n", + " %2 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo, #stablehlo]} : (tensor<1x128x128x1xf32>, tensor<3x3x1x1xf32>) -> tensor<1x128x128x1xf32>\n", " %3 = stablehlo.multiply %2, %2 : tensor<1x128x128x1xf32>\n", - " %4 = stablehlo.convolution(%arg0, %1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x128x128x1xf32>, tensor<3x3x1x1xf32>) -> tensor<1x128x128x1xf32>\n", + " %4 = stablehlo.convolution(%arg0, %1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo, #stablehlo]} : (tensor<1x128x128x1xf32>, tensor<3x3x1x1xf32>) -> tensor<1x128x128x1xf32>\n", " %5 = stablehlo.multiply %4, %4 : tensor<1x128x128x1xf32>\n", " %6 = stablehlo.add %3, %5 : tensor<1x128x128x1xf32>\n", " %7 = stablehlo.sqrt %6 : tensor<1x128x128x1xf32>\n", " return %7 : tensor<1x128x128x1xf32>\n", " }\n", - "}\n", - "\n", - "Wrote MLIR to path '/tmp/iree/colab_artifacts/edge_detection.mlir'\n" + "}" ] } ] @@ -369,7 +371,11 @@ { "cell_type": "code", "metadata": { - "id": "Ytvb5Gx_EFJl" + "id": "Ytvb5Gx_EFJl", + "outputId": "b9f8452f-1eeb-4355-8b2c-09d71dc38ec5", + "colab": { + "base_uri": "https://localhost:8080/" + } }, "source": [ "#@title Compile and prepare to test the edge detection module\n", @@ -383,7 +389,16 @@ "ctx.add_vm_module(vm_module)" ], "execution_count": 12, - "outputs": [] + "outputs": [ + { + "output_type": "stream", + "name": "stderr", + "text": [ + ":8: UserWarning: Making copy of unaligned VmModule buffer. It is recommended to make this deterministic by calling `copy_buffer` to always make a copy or `mmap` to efficiently load from a file. This warning can be silenced by adding `warn_if_copy=False` to `from_buffer`\n", + " vm_module = ireert.VmModule.from_flatbuffer(ctx.instance, flatbuffer_blob)\n" + ] + } + ] }, { "cell_type": "code", @@ -393,7 +408,7 @@ "height": 356 }, "id": "OUUXxol7wl-f", - "outputId": "b6b5a2b9-89b6-4574-e7aa-49803c28a2a4" + "outputId": "32b177e5-f3c1-4389-bf2e-69b4beaa5865" }, "source": [ "edge_detect_sobel_operator_f = ctx.modules.module[\"edge_detect_sobel_operator\"]\n", diff --git a/samples/colab/mnist_training.ipynb b/samples/colab/mnist_training.ipynb index 9f3a3ef0aca7..f7d3fab65c9c 100644 --- a/samples/colab/mnist_training.ipynb +++ b/samples/colab/mnist_training.ipynb @@ -22,19 +22,18 @@ } }, { - "cell_type": "markdown", + "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": { - "id": "license" + "cellView": "form", + "id": "dfJ7DaWpuMfS" }, - "source": [ - "```\n", - "Copyright 2020 The IREE Authors\n", - "\n", - "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\n", - "```\n" - ] + "execution_count": 1, + "outputs": [] }, { "cell_type": "markdown", @@ -62,7 +61,45 @@ "id": "Bn2JUIHS0P_Q" }, "source": [ - "## 1. Install and Import IREE" + "## 1. Install dependencies" + ] + }, + { + "cell_type": "code", + "source": [ + "%%capture\n", + "!python -m pip install --upgrade tf-nightly # Needed for stablehlo export in TF>=2.14" + ], + "metadata": { + "id": "NrpK8BOQuywf" + }, + "execution_count": 2, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "import tensorflow as tf\n", + "\n", + "# Print version information for future notebook users to reference.\n", + "print(\"TensorFlow version: \", tf.__version__)" + ], + "metadata": { + "id": "9ynEl1BMu3JI", + "outputId": "f54754f9-e3c1-4855-f984-0f0863469047", + "colab": { + "base_uri": "https://localhost:8080/" + } + }, + "execution_count": 3, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "TensorFlow version: 2.15.0-dev20230831\n" + ] + } ] }, { @@ -74,21 +111,39 @@ "%%capture\n", "!python -m pip install iree-compiler iree-runtime iree-tools-tf -f https://openxla.github.io/iree/pip-release-links.html" ], - "execution_count": 1, + "execution_count": 4, "outputs": [] }, { "cell_type": "code", "metadata": { - "id": "D0bOS2B50bL3" + "id": "D0bOS2B50bL3", + "outputId": "fe40d20f-26b1-47fc-9476-abe96a5d6269", + "colab": { + "base_uri": "https://localhost:8080/" + } }, "source": [ "# Import IREE's TensorFlow Compiler and Runtime.\n", "import iree.compiler.tf\n", - "import iree.runtime" + "import iree.runtime\n", + "\n", + "# Print version information for future notebook users to reference.\n", + "!iree-compile --version" ], - "execution_count": 2, - "outputs": [] + "execution_count": 5, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "IREE (https://openxla.github.io/iree):\n", + " IREE compiler version 20230831.630 @ 9ed3dab7ac4fcda959f5b8ebbcd7732aeb4b0c8d\n", + " LLVM version 18.0.0git\n", + " Optimized build\n" + ] + } + ] }, { "cell_type": "markdown", @@ -106,7 +161,7 @@ "colab": { "base_uri": "https://localhost:8080/" }, - "outputId": "5004eb6e-1f85-4257-eee7-739786e33b3c" + "outputId": "ec6138e1-1272-448d-ea26-83b5d54ca913" }, "source": [ "from matplotlib import pyplot as plt\n", @@ -125,21 +180,21 @@ "print(\"TensorFlow version: \", tf.__version__)\n", "print(\"Numpy version: \", np.__version__)" ], - "execution_count": 3, + "execution_count": 6, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ - "TensorFlow version: 2.12.0\n", - "Numpy version: 1.22.4\n" + "TensorFlow version: 2.15.0-dev20230831\n", + "Numpy version: 1.23.5\n" ] }, { "output_type": "stream", "name": "stderr", "text": [ - ":8: MatplotlibDeprecationWarning: The seaborn styles shipped by Matplotlib are deprecated since 3.6, as they no longer correspond to the styles shipped by seaborn. However, they will remain available as 'seaborn-v0_8-