samples/colab/tensorflow_hub_import.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {},
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "##### Copyright 2021 The IREE Authors"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "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"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "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://iree.dev/guides/ml-frameworks/tensorflow/."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Setup"
      ]
    },
    {
      "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 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": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "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": {},
      "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": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "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": {},
      "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": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "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": {},
      "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": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "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": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Import and compile the SavedModel with IREE"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "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 an intermediate .mlir file for future inspection.\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_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 iree_input to      '{iree_input}'\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "#@title Download compilation artifacts\n",
        "\n",
        "ARTIFACTS_ZIP = \"/tmp/mobilenet_colab_artifacts.zip\"\n",
        "\n",
        "print(f\"Zipping '{ARTIFACTS_DIR}' to '{ARTIFACTS_ZIP}' for download...\")\n",
        "!cd {ARTIFACTS_DIR} && zip -r {ARTIFACTS_ZIP} .\n",
        "\n",
        "# Note: you can also download files using the file explorer on the left\n",
        "try:\n",
        "    from google.colab import files\n",
        "    print(\"Downloading the artifacts zip file...\")\n",
        "    files.download(ARTIFACTS_ZIP)\n",
        "except ImportError:\n",
        "    print(\"Missing google_colab Python package, can't download files\")"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}