{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "tensorflow_hub_import.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "-V0X0E7LkEa4"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-V0X0E7LkEa4"
      },
      "source": [
        "##### Copyright 2021 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QEnB8SrBjubr"
      },
      "source": [
        "#@title License header\n",
        "# Copyright 2021 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": 1,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qb3S0mSjpK7J"
      },
      "source": [
        "# IREE TensorFlow Hub Import\n",
        "\n",
        "This notebook demonstrates how to download, import, and compile models from [TensorFlow Hub](https://tfhub.dev/). It covers:\n",
        "\n",
        "* Downloading a model from TensorFlow Hub\n",
        "* Ensuring the model has serving signatures needed for import\n",
        "* Importing and compiling the model with IREE\n",
        "\n",
        "At the end of the notebook, the compilation artifacts are compressed into a .zip file for you to download and use in an application.\n",
        "\n",
        "See also https://google.github.io/iree/ml-frameworks/tensorflow/."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9rNAJKNVkKOr"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RdVc4TbOkHM2"
      },
      "source": [
        "%%capture\n",
        "!python -m pip install iree-compiler-snapshot iree-runtime-snapshot iree-tools-tf-snapshot -f https://github.com/google/iree/releases"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qRwv3qI_l5O_",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "0016f7c2-c0f5-4339-d852-81019c4b8398"
      },
      "source": [
        "import os\n",
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub\n",
        "import tempfile\n",
        "from IPython.display import clear_output\n",
        "\n",
        "from iree.compiler import tf as tfc\n",
        "\n",
        "# Print version information for future notebook users to reference.\n",
        "print(\"TensorFlow version: \", tf.__version__)\n",
        "\n",
        "ARTIFACTS_DIR = os.path.join(tempfile.gettempdir(), \"iree\", \"colab_artifacts\")\n",
        "os.makedirs(ARTIFACTS_DIR, exist_ok=True)\n",
        "print(f\"Using artifacts directory '{ARTIFACTS_DIR}'\")"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "TensorFlow version:  2.4.1\n",
            "Using artifacts directory '/tmp/iree/colab_artifacts'\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZZAobcAhocFE"
      },
      "source": [
        "## Import pretrained [`mobilenet_v2`](https://tfhub.dev/google/tf2-preview/mobilenet_v2/classification/4) model\n",
        "\n",
        "IREE supports importing TensorFlow 2 models exported in the [SavedModel](https://www.tensorflow.org/guide/saved_model) format. This model we'll be importing is published in that format already, while other models may need to be converted first.\n",
        "\n",
        "MobileNet V2 is a family of neural network architectures for efficient on-device image classification and related tasks. This TensorFlow Hub module contains a trained instance of one particular network architecture packaged to perform image classification."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7fd0vmnloZo9",
        "outputId": "b5751219-0e88-4aa4-c4e2-e8645309e9e0"
      },
      "source": [
        "#@title Download the pretrained model\n",
        "\n",
        "# Use the `hub` library to download the pretrained model to the local disk\n",
        "# https://www.tensorflow.org/hub/api_docs/python/hub\n",
        "HUB_PATH = \"https://tfhub.dev/google/tf2-preview/mobilenet_v2/classification/4\"\n",
        "model_path = hub.resolve(HUB_PATH)\n",
        "print(f\"Downloaded model from tfhub to path: '{model_path}'\")"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Downloaded model from tfhub to path: '/tmp/tfhub_modules/426589ad685896ab7954855255a52db3442cb38d'\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CedNRSQTOE7C"
      },
      "source": [
        "### Check for serving signatures and re-export as needed\n",
        "\n",
        "IREE's compiler tools, like TensorFlow's `saved_model_cli` and other tools, require \"serving signatures\" to be defined in SavedModels.\n",
        "\n",
        "More references:\n",
        "\n",
        "* https://www.tensorflow.org/tfx/serving/signature_defs\n",
        "* https://blog.tensorflow.org/2021/03/a-tour-of-savedmodel-signatures.html"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qiO66oEYQmsd",
        "outputId": "b04179cf-f9dd-4624-a4df-db7190ec8e10"
      },
      "source": [
        "#@title Check for serving signatures\n",
        "\n",
        "# Load the SavedModel from the local disk and check if it has serving signatures\n",
        "# https://www.tensorflow.org/guide/saved_model#loading_and_using_a_custom_model\n",
        "loaded_model = tf.saved_model.load(model_path)\n",
        "serving_signatures = list(loaded_model.signatures.keys())\n",
        "print(f\"Loaded SavedModel from '{model_path}'\")\n",
        "print(f\"Serving signatures: {serving_signatures}\")\n",
        "\n",
        "# Also check with the saved_model_cli:\n",
        "print(\"\\n---\\n\")\n",
        "print(\"Checking for signature_defs using saved_model_cli:\\n\")\n",
        "!saved_model_cli show --dir {model_path} --tag_set serve --signature_def serving_default"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Loaded SavedModel from '/tmp/tfhub_modules/426589ad685896ab7954855255a52db3442cb38d'\n",
            "Serving signatures: []\n",
            "\n",
            "---\n",
            "\n",
            "Checking for signature_defs using saved_model_cli:\n",
            "\n",
            "The given SavedModel SignatureDef contains the following input(s):\n",
            "The given SavedModel SignatureDef contains the following output(s):\n",
            "Method name is: \n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kKqqX2LsReNz"
      },
      "source": [
        "Since the model we downloaded did not include any serving signatures, we'll re-export it with serving signatures defined.\n",
        "\n",
        "* https://www.tensorflow.org/guide/saved_model#specifying_signatures_during_export"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OlDG2OuqOBGC",
        "outputId": "61c77c94-3bc2-426c-c90c-78926a3d2cfa"
      },
      "source": [
        "#@title Look up input signatures to use when exporting\n",
        "\n",
        "# To save serving signatures we need to specify a `ConcreteFunction` with a\n",
        "# TensorSpec signature. We can determine what this signature should be by\n",
        "# looking at any documentation for the model or running the saved_model_cli.\n",
        "\n",
        "!saved_model_cli show --dir {model_path} --all \\\n",
        "    2> /dev/null | grep \"inputs: TensorSpec\" | tail -n 1"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "          inputs: TensorSpec(shape=(None, 224, 224, 3), dtype=tf.float32, name=u'inputs')\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gnb4HhMmkgiT",
        "outputId": "5278de21-0c74-4927-b8cb-21ca73fd43ed"
      },
      "source": [
        "#@title Re-export the model using the known signature\n",
        "\n",
        "# Get a concrete function using the signature we found above.\n",
        "# \n",
        "# The first element of the shape is a dynamic batch size. We'll be running\n",
        "# inference on a single image at a time, so set it to `1`. The rest of the\n",
        "# shape is the fixed image dimensions [width=224, height=224, channels=3].\n",
        "call = loaded_model.__call__.get_concrete_function(tf.TensorSpec([1, 224, 224, 3], tf.float32))\n",
        "\n",
        "# Save the model, setting the concrete function as a serving signature.\n",
        "# https://www.tensorflow.org/guide/saved_model#saving_a_custom_model\n",
        "resaved_model_path = '/tmp/resaved_model'\n",
        "tf.saved_model.save(loaded_model, resaved_model_path, signatures=call)\n",
        "clear_output()  # Skip over TensorFlow's output.\n",
        "print(f\"Saved model with serving signatures to '{resaved_model_path}'\")\n",
        "\n",
        "# Load the model back into memory and check that it has serving signatures now\n",
        "reloaded_model = tf.saved_model.load(resaved_model_path)\n",
        "reloaded_serving_signatures = list(reloaded_model.signatures.keys())\n",
        "print(f\"\\nReloaded SavedModel from '{resaved_model_path}'\")\n",
        "print(f\"Serving signatures: {reloaded_serving_signatures}\")\n",
        "\n",
        "# Also check with the saved_model_cli:\n",
        "print(\"\\n---\\n\")\n",
        "print(\"Checking for signature_defs using saved_model_cli:\\n\")\n",
        "!saved_model_cli show --dir {resaved_model_path} --tag_set serve --signature_def serving_default"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Saved model with serving signatures to '/tmp/resaved_model'\n",
            "\n",
            "Reloaded SavedModel from '/tmp/resaved_model'\n",
            "Serving signatures: ['serving_default']\n",
            "\n",
            "---\n",
            "\n",
            "Checking for signature_defs using saved_model_cli:\n",
            "\n",
            "The given SavedModel SignatureDef contains the following input(s):\n",
            "  inputs['inputs'] tensor_info:\n",
            "      dtype: DT_FLOAT\n",
            "      shape: (1, 224, 224, 3)\n",
            "      name: serving_default_inputs:0\n",
            "The given SavedModel SignatureDef contains the following output(s):\n",
            "  outputs['output_0'] tensor_info:\n",
            "      dtype: DT_FLOAT\n",
            "      shape: (1, 1001)\n",
            "      name: StatefulPartitionedCall:0\n",
            "Method name is: tensorflow/serving/predict\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YdmgASzwanSz"
      },
      "source": [
        "### Import and compile the SavedModel with IREE"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GLkjlHE5mdmg",
        "outputId": "4e25eabb-0355-4c77-f4c0-6c7a9f9ea8e3"
      },
      "source": [
        "#@title Import from SavedModel\n",
        "\n",
        "# The main output file from compilation is a .vmfb \"VM FlatBuffer\". This file\n",
        "# can used to run the compiled model with IREE's runtime.\n",
        "output_file = os.path.join(ARTIFACTS_DIR, \"mobilenet_v2.vmfb\")\n",
        "# As compilation runs, dump some intermediate .mlir files for future inspection.\n",
        "tf_input = os.path.join(ARTIFACTS_DIR, \"mobilenet_v2_tf_input.mlir\")\n",
        "iree_input = os.path.join(ARTIFACTS_DIR, \"mobilenet_v2_iree_input.mlir\")\n",
        "\n",
        "# Since our SavedModel uses signature defs, we use `saved_model_tags` with\n",
        "# `import_type=\"SIGNATURE_DEF\"`. If the SavedModel used an object graph, we\n",
        "# would use `exported_names` with `import_type=\"OBJECT_GRAPH\"` instead.\n",
        "\n",
        "# We'll set `target_backends=[\"vmvx\"]` to use IREE's reference CPU backend.\n",
        "# We could instead use different backends here, or set `import_only=True` then\n",
        "# download the imported .mlir file for compilation using native tools directly.\n",
        "\n",
        "tfc.compile_saved_model(\n",
        "    resaved_model_path,\n",
        "    output_file=output_file,\n",
        "    save_temp_tf_input=tf_input,\n",
        "    save_temp_iree_input=iree_input,\n",
        "    import_type=\"SIGNATURE_DEF\",\n",
        "    saved_model_tags=set([\"serve\"]),\n",
        "    target_backends=[\"vmvx\"])\n",
        "clear_output()  # Skip over TensorFlow's output.\n",
        "\n",
        "print(f\"Saved compiled output to '{output_file}'\")\n",
        "print(f\"Saved tf_input to        '{tf_input}'\")\n",
        "print(f\"Saved iree_input to      '{iree_input}'\")"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Saved compiled output to '/tmp/iree/colab_artifacts/mobilenet_v2.vmfb'\n",
            "Saved tf_input to        '/tmp/iree/colab_artifacts/mobilenet_v2_tf_input.mlir'\n",
            "Saved iree_input to      '/tmp/iree/colab_artifacts/mobilenet_v2_iree_input.mlir'\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 119
        },
        "id": "IEJAzOb5qASI",
        "outputId": "8d7c9b19-8891-43b2-de20-ebeeeafcff70"
      },
      "source": [
        "#@title Download compilation artifacts\n",
        "\n",
        "ARTIFACTS_ZIP = \"/tmp/mobilenet_colab_arficats.zip\"\n",
        "\n",
        "print(f\"Zipping '{ARTIFACTS_DIR}' to '{ARTIFACTS_ZIP}' for download...\")\n",
        "!zip -r {ARTIFACTS_ZIP} {ARTIFACTS_DIR}\n",
        "\n",
        "# Note: you can also download files using the file explorer on the left\n",
        "from google.colab import files\n",
        "print(\"Downloading the artifacts zip file...\")\n",
        "files.download(ARTIFACTS_ZIP)"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Zipping '/tmp/iree/colab_artifacts' to '/tmp/mobilenet_colab_arficats.zip' for download...\n",
            "  adding: tmp/iree/colab_artifacts/ (stored 0%)\n",
            "  adding: tmp/iree/colab_artifacts/mobilenet_v2_tf_input.mlir (deflated 47%)\n",
            "  adding: tmp/iree/colab_artifacts/mobilenet_v2.vmfb (deflated 8%)\n",
            "  adding: tmp/iree/colab_artifacts/mobilenet_v2_iree_input.mlir (deflated 47%)\n",
            "Downloading the artifacts zip file...\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/javascript": [
              "\n",
              "    async function download(id, filename, size) {\n",
              "      if (!google.colab.kernel.accessAllowed) {\n",
              "        return;\n",
              "      }\n",
              "      const div = document.createElement('div');\n",
              "      const label = document.createElement('label');\n",
              "      label.textContent = `Downloading \"${filename}\": `;\n",
              "      div.appendChild(label);\n",
              "      const progress = document.createElement('progress');\n",
              "      progress.max = size;\n",
              "      div.appendChild(progress);\n",
              "      document.body.appendChild(div);\n",
              "\n",
              "      const buffers = [];\n",
              "      let downloaded = 0;\n",
              "\n",
              "      const channel = await google.colab.kernel.comms.open(id);\n",
              "      // Send a message to notify the kernel that we're ready.\n",
              "      channel.send({})\n",
              "\n",
              "      for await (const message of channel.messages) {\n",
              "        // Send a message to notify the kernel that we're ready.\n",
              "        channel.send({})\n",
              "        if (message.buffers) {\n",
              "          for (const buffer of message.buffers) {\n",
              "            buffers.push(buffer);\n",
              "            downloaded += buffer.byteLength;\n",
              "            progress.value = downloaded;\n",
              "          }\n",
              "        }\n",
              "      }\n",
              "      const blob = new Blob(buffers, {type: 'application/binary'});\n",
              "      const a = document.createElement('a');\n",
              "      a.href = window.URL.createObjectURL(blob);\n",
              "      a.download = filename;\n",
              "      div.appendChild(a);\n",
              "      a.click();\n",
              "      div.remove();\n",
              "    }\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "application/javascript": [
              "download(\"download_23ce7642-800c-4e5e-afdb-7c6541fc9cd9\", \"mobilenet_colab_arficats.zip\", 43802015)"
            ],
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    }
  ]
}