TF-TRT-inference-from-saved-model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TF-TRT Inference From Saved Model with Tensorflow <= 1.13\n",
    "\n",
    "In this notebook, we demonstrate the process to create a TF-TRT optimized model from a Tensorflow *saved model*. This notebook has been successfully tested with the NVIDIA NGC Tensorflow 1.13 container `nvcr.io/nvidia/tensorflow:19.04-py3` that can be downloaded freely from http://ngc.nvidia.com.\n",
    "\n",
    "## Notebook  Content\n",
    "1. [Pre-requisite: data and model](#1)\n",
    "1. [Verifying the orignal FP32 model](#2)\n",
    "1. [Creating TF-TRT FP32 model](#3)\n",
    "1. [Creating TF-TRT FP16 model](#4)\n",
    "1. [Creating TF-TRT INT8 model](#5)\n",
    "1. [Calibrating TF-TRT INT8 model with raw JPEG images](#6)\n",
    " \n",
    "## Quick start\n",
    "We will run this demonstration with a saved Resnet-v1-50 model stored at `/path/to/saved_model`.\n",
    "\n",
    "The INT8 calibration process requires access to a small but representative sample of real training or valiation data.\n",
    "\n",
    "We will use the ImageNet dataset that is stored in TFrecords format. Google provide an excellent all-in-one script for downloading and preparing the ImageNet dataset at \n",
    "\n",
    "https://github.com/tensorflow/models/blob/master/research/inception/inception/data/download_and_preprocess_imagenet.sh.\n",
    "\n",
    "\n",
    "To run this notebook, start the NGC TF container, providing correct path to the ImageNet validation data `/path/to/image_net` and the folder `/path/to/saved_model` containing the TF saved model:\n",
    "\n",
    "```bash\n",
    "nvidia-docker run --rm -it -p 8888:8888 -v /path/to/image_net:/data  -v /path/to/saved_model:/saved_model --name TFTRT nvcr.io/nvidia/tensorflow:19.04-py3\n",
    "```\n",
    "\n",
    "Within the container, we then start Jupyter notebook with:\n",
    "\n",
    "```bash\n",
    "jupyter notebook --ip 0.0.0.0 --port 8888  --allow-root\n",
    "```\n",
    "\n",
    "Connect to Jupyter notebook web interface on your host http://localhost:8888.\n",
    "\n",
    "\n",
    "<a id=\"1\"></a>\n",
    "## 1. Pre-requisite: data and model\n",
    "\n",
    "We first install some extra packages and external dependencies needed for, e.g. preprocessing ImageNet data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "pushd /workspace/nvidia-examples/tensorrt/tftrt/examples/object_detection/ \n",
    "bash install_dependencies.sh;\n",
    "popd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ii  libnvinfer-dev                         5.1.2-1+cuda10.1                      amd64        TensorRT development libraries and headers\r\n",
      "ii  libnvinfer5                            5.1.2-1+cuda10.1                      amd64        TensorRT runtime libraries\r\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow.contrib.tensorrt as trt\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import time\n",
    "\n",
    "import logging\n",
    "logging.getLogger(\"tensorflow\").setLevel(logging.ERROR)\n",
    "\n",
    "import os\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\n",
    "\n",
    "config = tf.ConfigProto()\n",
    "config.gpu_options.allow_growth=True\n",
    "\n",
    "#check TensorRT version\n",
    "!dpkg -l | grep nvinfer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data\n",
    "We verify that the correct Imagenet data folder has been mounted and validation data files of the form `validation-00xxx-of-00128` are available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_files(data_dir, filename_pattern):\n",
    "    if data_dir == None:\n",
    "        return []\n",
    "    files = tf.gfile.Glob(os.path.join(data_dir, filename_pattern))\n",
    "    if files == []:\n",
    "        raise ValueError('Can not find any files in {} with '\n",
    "                         'pattern \"{}\"'.format(data_dir, filename_pattern))\n",
    "    return files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "There are 128 calibration files. \n",
      "/data/validation-00114-of-00128\n",
      "/data/validation-00094-of-00128\n",
      "...\n"
     ]
    }
   ],
   "source": [
    "VALIDATION_DATA_DIR = \"/data\"\n",
    "calibration_files = get_files(VALIDATION_DATA_DIR, 'validation*')\n",
    "print('There are %d calibration files. \\n%s\\n%s\\n...'%(len(calibration_files), calibration_files[0], calibration_files[-1]))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TF saved model\n",
    "If not already downloaded, we will be downloading and working with a ResNet-50 v1 checkpoint from https://github.com/tensorflow/models/tree/master/official/resnet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "FILE=/saved_model/resnet_v1_50_2016_08_28.tar.gz\n",
    "if [ -f $FILE ]; then\n",
    "   echo \"The file '$FILE' exists.\"\n",
    "else\n",
    "   echo \"The file '$FILE' in not found. Downloading...\"\n",
    "   wget -P /saved_model/ http://download.tensorflow.org/models/official/20181001_resnet/savedmodels/resnet_v1_fp32_savedmodel_NHWC.tar.gz\n",
    "fi\n",
    "\n",
    "tar -xzvf /saved_model/resnet_v1_fp32_savedmodel_NHWC.tar.gz -C /saved_model "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Helper functions\n",
    "We define a few helper functions to read and preprocess Imagenet data from TFRecord files. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def deserialize_image_record(record):\n",
    "    feature_map = {\n",
    "        'image/encoded':          tf.FixedLenFeature([ ], tf.string, ''),\n",
    "        'image/class/label':      tf.FixedLenFeature([1], tf.int64,  -1),\n",
    "        'image/class/text':       tf.FixedLenFeature([ ], tf.string, ''),\n",
    "        'image/object/bbox/xmin': tf.VarLenFeature(dtype=tf.float32),\n",
    "        'image/object/bbox/ymin': tf.VarLenFeature(dtype=tf.float32),\n",
    "        'image/object/bbox/xmax': tf.VarLenFeature(dtype=tf.float32),\n",
    "        'image/object/bbox/ymax': tf.VarLenFeature(dtype=tf.float32)\n",
    "    }\n",
    "    with tf.name_scope('deserialize_image_record'):\n",
    "        obj = tf.parse_single_example(record, feature_map)\n",
    "        imgdata = obj['image/encoded']\n",
    "        label   = tf.cast(obj['image/class/label'], tf.int32)\n",
    "        bbox    = tf.stack([obj['image/object/bbox/%s'%x].values\n",
    "                            for x in ['ymin', 'xmin', 'ymax', 'xmax']])\n",
    "        bbox = tf.transpose(tf.expand_dims(bbox, 0), [0,2,1])\n",
    "        text    = obj['image/class/text']\n",
    "        return imgdata, label, bbox, text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from preprocessing import vgg_preprocessing\n",
    "def preprocess(record):\n",
    "        # Parse TFRecord\n",
    "        imgdata, label, bbox, text = deserialize_image_record(record)\n",
    "        #label -= 1 # Change to 0-based if not using background class\n",
    "        try:    image = tf.image.decode_jpeg(imgdata, channels=3, fancy_upscaling=False, dct_method='INTEGER_FAST')\n",
    "        except: image = tf.image.decode_png(imgdata, channels=3)\n",
    "\n",
    "        image = vgg_preprocessing.preprocess_image(image, 224, 224, is_training=False)\n",
    "        return image, label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Define some global variables\n",
    "BATCH_SIZE = 64\n",
    "\n",
    "dataset = tf.data.TFRecordDataset(calibration_files)    \n",
    "dataset = dataset.apply(tf.contrib.data.map_and_batch(map_func=preprocess, batch_size=BATCH_SIZE, num_parallel_calls=20))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"2\"></a>\n",
    "## 2. Verifying the orignal FP32 model\n",
    "We demonstrate the conversion process with a Resnet-50 v1 model. First, we inspect the original Tensorflow model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "SAVED_MODEL_DIR =  \"/saved_model/resnet_v1_fp32_savedmodel_NHWC/1538686669/\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We employ `saved_model_cli` to inspect the inputs and outputs of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:\n",
      "\n",
      "signature_def['predict']:\n",
      "  The given SavedModel SignatureDef contains the following input(s):\n",
      "    inputs['input'] tensor_info:\n",
      "        dtype: DT_FLOAT\n",
      "        shape: (64, 224, 224, 3)\n",
      "        name: input_tensor:0\n",
      "  The given SavedModel SignatureDef contains the following output(s):\n",
      "    outputs['classes'] tensor_info:\n",
      "        dtype: DT_INT64\n",
      "        shape: (64)\n",
      "        name: ArgMax:0\n",
      "    outputs['probabilities'] tensor_info:\n",
      "        dtype: DT_FLOAT\n",
      "        shape: (64, 1001)\n",
      "        name: softmax_tensor:0\n",
      "  Method name is: tensorflow/serving/predict\n",
      "\n",
      "signature_def['serving_default']:\n",
      "  The given SavedModel SignatureDef contains the following input(s):\n",
      "    inputs['input'] tensor_info:\n",
      "        dtype: DT_FLOAT\n",
      "        shape: (64, 224, 224, 3)\n",
      "        name: input_tensor:0\n",
      "  The given SavedModel SignatureDef contains the following output(s):\n",
      "    outputs['classes'] tensor_info:\n",
      "        dtype: DT_INT64\n",
      "        shape: (64)\n",
      "        name: ArgMax:0\n",
      "    outputs['probabilities'] tensor_info:\n",
      "        dtype: DT_FLOAT\n",
      "        shape: (64, 1001)\n",
      "        name: softmax_tensor:0\n",
      "  Method name is: tensorflow/serving/predict\n"
     ]
    }
   ],
   "source": [
    "!saved_model_cli show --all --dir $SAVED_MODEL_DIR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This give us information on the input and output tensors as `input_tensor:0` and `softmax_tensor:0` respectively. Also note that the number of output classes here is 1001 instead of 1000 Imagenet classes. This is because the network was trained with an extra background class. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "INPUT_TENSOR = 'input_tensor:0'\n",
    "OUTPUT_TENSOR = 'softmax_tensor:0'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we define a function to read in a saved mode, measuring its speed and accuracy on the validation data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def benchmark_saved_model(SAVED_MODEL_DIR, dataset=dataset, BATCH_SIZE=64):\n",
    "    with tf.Session(graph=tf.Graph(), config=config) as sess:\n",
    "        # prepare dataset iterator\n",
    "        iterator = dataset.make_one_shot_iterator()\n",
    "        next_element = iterator.get_next()\n",
    "\n",
    "        tf.saved_model.loader.load(\n",
    "            sess, [tf.saved_model.tag_constants.SERVING], SAVED_MODEL_DIR)\n",
    "\n",
    "        print('Warming up for 50 batches...')\n",
    "        for _ in range (50):\n",
    "            sess.run(OUTPUT_TENSOR, feed_dict={INPUT_TENSOR: sess.run(next_element)[0]})\n",
    "\n",
    "        print('Benchmarking inference engine...')\n",
    "        num_hits = 0\n",
    "        num_predict = 0\n",
    "        start_time = time.time()\n",
    "        try:\n",
    "            while True:        \n",
    "                image_data = sess.run(next_element)    \n",
    "                img = image_data[0]\n",
    "                label = image_data[1].squeeze()\n",
    "                output = sess.run([OUTPUT_TENSOR], feed_dict={INPUT_TENSOR: img})            \n",
    "                prediction = np.argmax(output[0], axis=1)\n",
    "                num_hits += np.sum(prediction == label)\n",
    "                num_predict += len(prediction)\n",
    "        except tf.errors.OutOfRangeError as e:\n",
    "            pass\n",
    "\n",
    "        print('Accuracy: %.2f%%'%(100*num_hits/num_predict))\n",
    "        print('Inference speed: %.2f samples/s'%(num_predict/(time.time()-start_time)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "benchmark_saved_model(SAVED_MODEL_DIR, dataset=dataset, BATCH_SIZE=BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"3\"></a>\n",
    "## 3. Creating TF-TRT FP32 model\n",
    "\n",
    "Next, we convert the naitive TF FP32 model to TF-TRT FP32, then verify model accuracy and inference speed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "FP32_SAVED_MODEL_DIR = SAVED_MODEL_DIR+\"_TFTRT_FP32/1\"\n",
    "!rm -rf $FP32_SAVED_MODEL_DIR\n",
    "\n",
    "#Now we create the TFTRT FP32 engine\n",
    "trt.create_inference_graph(\n",
    "    input_graph_def=None,\n",
    "    outputs=None,\n",
    "    max_batch_size=BATCH_SIZE,\n",
    "    input_saved_model_dir=SAVED_MODEL_DIR,\n",
    "    output_saved_model_dir=FP32_SAVED_MODEL_DIR,\n",
    "    precision_mode=\"FP32\")\n",
    "\n",
    "benchmark_saved_model(FP32_SAVED_MODEL_DIR, dataset=dataset, BATCH_SIZE=BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"4\"></a>\n",
    "## 4. Creating TF-TRT FP16 model\n",
    "\n",
    "Next, we convert the naitive TF FP32 model to TF-TRT FP16, then verify model accuracy and inference speed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "FP16_SAVED_MODEL_DIR = SAVED_MODEL_DIR+\"_TFTRT_FP16/1\"\n",
    "!rm -rf $FP16_SAVED_MODEL_DIR\n",
    "\n",
    "#Now we create the TFTRT FP32 engine\n",
    "trt.create_inference_graph(\n",
    "    input_graph_def=None,\n",
    "    outputs=None,\n",
    "    max_batch_size=BATCH_SIZE,\n",
    "    input_saved_model_dir=SAVED_MODEL_DIR,\n",
    "    output_saved_model_dir=FP16_SAVED_MODEL_DIR,\n",
    "    precision_mode=\"FP16\")\n",
    "\n",
    "benchmark_saved_model(FP16_SAVED_MODEL_DIR, dataset=dataset, BATCH_SIZE=BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"5\"></a>\n",
    "## 5. Creating TF-TRT INT8 model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating TF-TRT INT8 inference model requires two steps:\n",
    "\n",
    "- Step 1: creating the calibration graph, and run some training data through that graph for INT-8 calibration.\n",
    "\n",
    "- Step 2: converting the calibration graph to the TF-TRT INT8 inference engine\n",
    "\n",
    "### Step 1: Creating the calibration graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Preparing calibration data...\n",
      "Loading INT8 calibration graph...\n",
      "Calibrate model on calibration data...\n"
     ]
    }
   ],
   "source": [
    "#Now we create the TFTRT INT8 calibration graph\n",
    "trt_int8_calib_graph = trt.create_inference_graph(\n",
    "    input_graph_def=None,\n",
    "    outputs=[OUTPUT_TENSOR],\n",
    "    max_batch_size=BATCH_SIZE,\n",
    "    input_saved_model_dir=SAVED_MODEL_DIR,    \n",
    "    precision_mode=\"INT8\")\n",
    "\n",
    "#Then calibrate it with 10 batches of examples\n",
    "N_runs=10\n",
    "with tf.Session(graph=tf.Graph(), config=config) as sess:\n",
    "    print('Preparing calibration data...')\n",
    "    iterator = dataset.make_one_shot_iterator()\n",
    "    next_element = iterator.get_next()\n",
    "\n",
    "    print('Loading INT8 calibration graph...')\n",
    "    output_node = tf.import_graph_def(\n",
    "        trt_int8_calib_graph,\n",
    "        return_elements=[OUTPUT_TENSOR],\n",
    "        name='')\n",
    "\n",
    "    print('Calibrate model on calibration data...')    \n",
    "    for _ in range(N_runs):\n",
    "        sess.run(output_node, feed_dict={INPUT_TENSOR: sess.run(next_element)[0]})            "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Converting the calibration graph to inference graph\n",
    "\n",
    "Now we convert the INT8 calibration graph to the final TF-TRT INT8 inference engine, then save this engine to a *saved model*, ready to be served elsewhere."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " TF-TRT INT8 inference engine...\n"
     ]
    }
   ],
   "source": [
    "#Create Int8 inference model from the calibration graph and write to a saved session\n",
    "print('Creating TF-TRT INT8 inference engine...')\n",
    "trt_int8_calibrated_graph=trt.calib_graph_to_infer_graph(trt_int8_calib_graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading TF-TRT INT8 inference engine...\n",
      "Saving INT8 model to /saved_model/resnet_v1_fp32_savedmodel_NHWC/1538686669/_TFTRT_INT8/1\n"
     ]
    }
   ],
   "source": [
    "#set a directory to write the saved model\n",
    "INT8_SAVED_MODEL_DIR =  SAVED_MODEL_DIR + \"_TFTRT_INT8/1\"\n",
    "!rm -rf $INT8_SAVED_MODEL_DIR\n",
    "\n",
    "with tf.Session(graph=tf.Graph()) as sess:\n",
    "    print('Loading TF-TRT INT8 inference engine...')\n",
    "    output_node = tf.import_graph_def(\n",
    "        trt_int8_calibrated_graph,\n",
    "        return_elements=[OUTPUT_TENSOR],\n",
    "        name='')\n",
    "\n",
    "    #Save model for serving\n",
    "    print('Saving INT8 model to %s'%INT8_SAVED_MODEL_DIR)\n",
    "    tf.saved_model.simple_save(\n",
    "        session=sess,\n",
    "        export_dir=INT8_SAVED_MODEL_DIR,\n",
    "        inputs={\"input\":tf.get_default_graph().get_tensor_by_name(INPUT_TENSOR)},\n",
    "        outputs={\"softmax\":tf.get_default_graph().get_tensor_by_name(OUTPUT_TENSOR),\n",
    "                 \"classes\":tf.get_default_graph().get_tensor_by_name(\"ArgMax:0\")},\n",
    "        legacy_init_op=None\n",
    "     )    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Benchmarking INT8 saved model\n",
    "\n",
    "Finally we reload and verify the accuracy and performance of the INT8 saved model from disk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "benchmark_saved_model(INT8_SAVED_MODEL_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:\n",
      "\n",
      "signature_def['serving_default']:\n",
      "  The given SavedModel SignatureDef contains the following input(s):\n",
      "    inputs['input'] tensor_info:\n",
      "        dtype: DT_FLOAT\n",
      "        shape: (64, 224, 224, 3)\n",
      "        name: input_tensor:0\n",
      "  The given SavedModel SignatureDef contains the following output(s):\n",
      "    outputs['classes'] tensor_info:\n",
      "        dtype: DT_INT64\n",
      "        shape: unknown_rank\n",
      "        name: ArgMax:0\n",
      "    outputs['softmax'] tensor_info:\n",
      "        dtype: DT_FLOAT\n",
      "        shape: unknown_rank\n",
      "        name: softmax_tensor:0\n",
      "  Method name is: tensorflow/serving/predict\n"
     ]
    }
   ],
   "source": [
    "!saved_model_cli show --all --dir $INT8_SAVED_MODEL_DIR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"6\"></a>\n",
    "## 6. Calibrating TF-TRT INT8 model with raw JPEG images\n",
    "\n",
    "As an alternative to taking data in TFRecords format, in this section, we demonstrate the process of calibrating TFTRT INT-8 model from a directory of raw JPEG images. We asume that raw images have been mounted to the directory `/data/Calibration_data`.\n",
    "\n",
    "As a rule of thumb, calibration data should be a small but representative set of images that is similar to what is expected in deployment. Empirically, for common network architectures trained on imagenet data, calibration data of size 500-1000 provide good accuracy. As such, a good strategy for a dataset such as imagenet is to choose one sample from each class. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "There are 1000 calibration files. \n",
      "/data/Calibration_data/n03710193_29586.JPEG\n",
      "/data/Calibration_data/n02085782_2655.JPEG\n",
      "...\n"
     ]
    }
   ],
   "source": [
    "data_directory = \"/data/Calibration_data\"\n",
    "calibration_files = [os.path.join(path, name) for path, _, files in os.walk(data_directory) for name in files]\n",
    "print('There are %d calibration files. \\n%s\\n%s\\n...'%(len(calibration_files), calibration_files[0], calibration_files[-1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define a helper function to read and preprocess image from JPEG file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_file(filepath):\n",
    "    image = tf.read_file(filepath)\n",
    "    image = tf.image.decode_jpeg(image, channels=3)\n",
    "    image = vgg_preprocessing.preprocess_image(image, 224, 224, is_training=False)\n",
    "    return image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "dataset = tf.data.Dataset.from_tensor_slices(data_files)\n",
    "dataset = dataset.apply(tf.contrib.data.map_and_batch(map_func=parse_file, batch_size=BATCH_SIZE, num_parallel_calls=20))\n",
    "dataset = dataset.repeat(count=1)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we proceed with the two-stage process of creating and calibrating TFTRT INT8 model.\n",
    "\n",
    "### Step 1: Creating the calibration graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Preparing calibration data...\n",
      "Loading INT8 calibration graph...\n",
      "Calibrate model on calibration data...\n"
     ]
    }
   ],
   "source": [
    "#Now we create the TFTRT INT8 calibration graph\n",
    "trt_int8_calib_graph = trt.create_inference_graph(\n",
    "    input_graph_def=None,\n",
    "    outputs=[OUTPUT_TENSOR],\n",
    "    max_batch_size=BATCH_SIZE,\n",
    "    input_saved_model_dir=SAVED_MODEL_DIR,    \n",
    "    precision_mode=\"INT8\")\n",
    "\n",
    "#Then calibrate it with 10 batches of examples\n",
    "N_runs=10\n",
    "with tf.Session(graph=tf.Graph(), config=config) as sess:\n",
    "    print('Preparing calibration data...')\n",
    "    iterator = dataset.make_one_shot_iterator()\n",
    "    next_element = iterator.get_next()\n",
    "\n",
    "    print('Loading INT8 calibration graph...')\n",
    "    output_node = tf.import_graph_def(\n",
    "        trt_int8_calib_graph,\n",
    "        return_elements=[OUTPUT_TENSOR],\n",
    "        name='')\n",
    "\n",
    "    print('Calibrate model on calibration data...')    \n",
    "    for _ in range(N_runs):\n",
    "        sess.run(output_node, feed_dict={INPUT_TENSOR: sess.run(next_element)})            "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Converting the calibration graph to inference graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating TF-TRT INT8 inference engine...\n",
      "Loading TF-TRT INT8 inference engine...\n",
      "Saving INT8 model to /saved_model/resnet_v1_fp32_savedmodel_NHWC/1538686669/_TFTRT_INT8_JPEG/1\n"
     ]
    }
   ],
   "source": [
    "#Create Int8 inference model from the calibration graph and write to a saved session\n",
    "print('Creating TF-TRT INT8 inference engine...')\n",
    "trt_int8_calibrated_graph=trt.calib_graph_to_infer_graph(trt_int8_calib_graph)\n",
    "\n",
    "#set a directory to write the saved model\n",
    "INT8_SAVED_MODEL_DIR =  SAVED_MODEL_DIR + \"_TFTRT_INT8_JPEG/1\"\n",
    "!rm -rf $INT8_SAVED_MODEL_DIR\n",
    "\n",
    "with tf.Session(graph=tf.Graph()) as sess:\n",
    "    print('Loading TF-TRT INT8 inference engine...')\n",
    "    output_node = tf.import_graph_def(\n",
    "        trt_int8_calibrated_graph,\n",
    "        return_elements=[OUTPUT_TENSOR],\n",
    "        name='')\n",
    "\n",
    "    #Save model for serving\n",
    "    print('Saving INT8 model to %s'%INT8_SAVED_MODEL_DIR)\n",
    "    tf.saved_model.simple_save(\n",
    "        session=sess,\n",
    "        export_dir=INT8_SAVED_MODEL_DIR,\n",
    "        inputs={\"input\":tf.get_default_graph().get_tensor_by_name(INPUT_TENSOR)},\n",
    "        outputs={\"softmax\":tf.get_default_graph().get_tensor_by_name(OUTPUT_TENSOR),\n",
    "                 \"classes\":tf.get_default_graph().get_tensor_by_name(\"ArgMax:0\")},\n",
    "        legacy_init_op=None\n",
    "     )    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As before, we can benchmark the speed and accuracy of the resulting model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "benchmark_saved_model(INT8_SAVED_MODEL_DIR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "In this notebook, we have demonstrated the process of creating TF-TRT inference model from an original TF FP32 *saved model*. In every case, we have also verified the accuracy and speed to the resulting model. \n"
   ]
  }
 ],
 "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}