{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Compile Tensorflow Models\nThis article is an introductory tutorial to deploy tensorflow models with TVM.\n\nFor us to begin with, tensorflow python module is required to be installed.\n\nPlease refer to https://www.tensorflow.org/install\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# tvm, relay\nimport tvm\nfrom tvm import te\nfrom tvm import relay\n\n# os and numpy\nimport numpy as np\nimport os.path\n\n# Tensorflow imports\nimport tensorflow as tf\n\n\n# Ask tensorflow to limit its GPU memory to what's actually needed\n# instead of gobbling everything that's available.\n# https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth\n# This way this tutorial is a little more friendly to sphinx-gallery.\ngpus = tf.config.list_physical_devices(\"GPU\")\nif gpus:\n    try:\n        for gpu in gpus:\n            tf.config.experimental.set_memory_growth(gpu, True)\n        print(\"tensorflow will use experimental.set_memory_growth(True)\")\n    except RuntimeError as e:\n        print(\"experimental.set_memory_growth option is not available: {}\".format(e))\n\n\ntry:\n    tf_compat_v1 = tf.compat.v1\nexcept ImportError:\n    tf_compat_v1 = tf\n\n# Tensorflow utility functions\nimport tvm.relay.testing.tf as tf_testing\n\n# Base location for model related files.\nrepo_base = \"https://github.com/dmlc/web-data/raw/main/tensorflow/models/InceptionV1/\"\n\n# Test image\nimg_name = \"elephant-299.jpg\"\nimage_url = os.path.join(repo_base, img_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Tutorials\nPlease refer docs/frontend/tensorflow.md for more details for various models\nfrom tensorflow.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "model_name = \"classify_image_graph_def-with_shapes.pb\"\nmodel_url = os.path.join(repo_base, model_name)\n\n# Image label map\nmap_proto = \"imagenet_2012_challenge_label_map_proto.pbtxt\"\nmap_proto_url = os.path.join(repo_base, map_proto)\n\n# Human readable text for labels\nlabel_map = \"imagenet_synset_to_human_label_map.txt\"\nlabel_map_url = os.path.join(repo_base, label_map)\n\n# Target settings\n# Use these commented settings to build for cuda.\n# target = tvm.target.Target(\"cuda\", host=\"llvm\")\n# layout = \"NCHW\"\n# dev = tvm.cuda(0)\ntarget = tvm.target.Target(\"llvm\", host=\"llvm\")\nlayout = None\ndev = tvm.cpu(0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Download required files\nDownload files listed above.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from tvm.contrib.download import download_testdata\n\nimg_path = download_testdata(image_url, img_name, module=\"data\")\nmodel_path = download_testdata(model_url, model_name, module=[\"tf\", \"InceptionV1\"])\nmap_proto_path = download_testdata(map_proto_url, map_proto, module=\"data\")\nlabel_path = download_testdata(label_map_url, label_map, module=\"data\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Import model\nCreates tensorflow graph definition from protobuf file.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "with tf_compat_v1.gfile.GFile(model_path, \"rb\") as f:\n    graph_def = tf_compat_v1.GraphDef()\n    graph_def.ParseFromString(f.read())\n    graph = tf.import_graph_def(graph_def, name=\"\")\n    # Call the utility to import the graph definition into default graph.\n    graph_def = tf_testing.ProcessGraphDefParam(graph_def)\n    # Add shapes to the graph.\n    with tf_compat_v1.Session() as sess:\n        graph_def = tf_testing.AddShapesToGraphDef(sess, \"softmax\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Decode image\n<div class=\"alert alert-info\"><h4>Note</h4><p>tensorflow frontend import doesn't support preprocessing ops like JpegDecode.\n  JpegDecode is bypassed (just return source node).\n  Hence we supply decoded frame to TVM instead.</p></div>\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from PIL import Image\n\nimage = Image.open(img_path).resize((299, 299))\n\nx = np.array(image)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Import the graph to Relay\nImport tensorflow graph definition to relay frontend.\n\nResults:\n  sym: relay expr for given tensorflow protobuf.\n  params: params converted from tensorflow params (tensor protobuf).\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "shape_dict = {\"DecodeJpeg/contents\": x.shape}\ndtype_dict = {\"DecodeJpeg/contents\": \"uint8\"}\nmod, params = relay.frontend.from_tensorflow(graph_def, layout=layout, shape=shape_dict)\n\nprint(\"Tensorflow protobuf imported to relay frontend.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Relay Build\nCompile the graph to llvm target with given input specification.\n\nResults:\n  graph: Final graph after compilation.\n  params: final params after compilation.\n  lib: target library which can be deployed on target with TVM runtime.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "with tvm.transform.PassContext(opt_level=3):\n    lib = relay.build(mod, target, params=params)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Execute the portable graph on TVM\nNow we can try deploying the compiled model on target.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from tvm.contrib import graph_executor\n\ndtype = \"uint8\"\nm = graph_executor.GraphModule(lib[\"default\"](dev))\n# set inputs\nm.set_input(\"DecodeJpeg/contents\", tvm.nd.array(x.astype(dtype)))\n# execute\nm.run()\n# get outputs\ntvm_output = m.get_output(0, tvm.nd.empty(((1, 1008)), \"float32\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Process the output\nProcess the model output to human readable text for InceptionV1.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "predictions = tvm_output.numpy()\npredictions = np.squeeze(predictions)\n\n# Creates node ID --> English string lookup.\nnode_lookup = tf_testing.NodeLookup(label_lookup_path=map_proto_path, uid_lookup_path=label_path)\n\n# Print top 5 predictions from TVM output.\ntop_k = predictions.argsort()[-5:][::-1]\nfor node_id in top_k:\n    human_string = node_lookup.id_to_string(node_id)\n    score = predictions[node_id]\n    print(\"%s (score = %.5f)\" % (human_string, score))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Inference on tensorflow\nRun the corresponding model on tensorflow\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def create_graph():\n    \"\"\"Creates a graph from saved GraphDef file and returns a saver.\"\"\"\n    # Creates graph from saved graph_def.pb.\n    with tf_compat_v1.gfile.GFile(model_path, \"rb\") as f:\n        graph_def = tf_compat_v1.GraphDef()\n        graph_def.ParseFromString(f.read())\n        graph = tf.import_graph_def(graph_def, name=\"\")\n        # Call the utility to import the graph definition into default graph.\n        graph_def = tf_testing.ProcessGraphDefParam(graph_def)\n\n\ndef run_inference_on_image(image):\n    \"\"\"Runs inference on an image.\n\n    Parameters\n    ----------\n    image: String\n        Image file name.\n\n    Returns\n    -------\n        Nothing\n    \"\"\"\n    if not tf_compat_v1.gfile.Exists(image):\n        tf.logging.fatal(\"File does not exist %s\", image)\n    image_data = tf_compat_v1.gfile.GFile(image, \"rb\").read()\n\n    # Creates graph from saved GraphDef.\n    create_graph()\n\n    with tf_compat_v1.Session() as sess:\n        softmax_tensor = sess.graph.get_tensor_by_name(\"softmax:0\")\n        predictions = sess.run(softmax_tensor, {\"DecodeJpeg/contents:0\": image_data})\n\n        predictions = np.squeeze(predictions)\n\n        # Creates node ID --> English string lookup.\n        node_lookup = tf_testing.NodeLookup(\n            label_lookup_path=map_proto_path, uid_lookup_path=label_path\n        )\n\n        # Print top 5 predictions from tensorflow.\n        top_k = predictions.argsort()[-5:][::-1]\n        print(\"===== TENSORFLOW RESULTS =======\")\n        for node_id in top_k:\n            human_string = node_lookup.id_to_string(node_id)\n            score = predictions[node_id]\n            print(\"%s (score = %.5f)\" % (human_string, score))\n\n\nrun_inference_on_image(img_path)"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.5"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}