TFTRT-int8-calibration-from-saved-model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tensorflow - TensorRT INT8 Inference example from saved model\n",
    "\n",
    "In this notebook, we demonstrate the process to create a TF-TensorRT optimized model from a Tensorflow *saved model*.\n",
    "This notebook has been successfully tested in the NVIDIA NGC Tensorflow container `nvcr.io/nvidia/tensorflow:19.04-py3` that can be downloaded from http://ngc.nvidia.com.\n",
    "\n",
    "### Data\n",
    "We use the ImageNet dataset that has been stored in TFrecords format. Google provide an excellent all-in-one script for downloading and preparing the ImageNet dataset at https://github.com/tensorflow/models/blob/master/research/inception/inception/data/download_and_preprocess_imagenet.sh.\n",
    "\n",
    "### Saved model\n",
    "We will run this demonstration with a saved model from the Tensorflow Resnet model zoo https://github.com/tensorflow/models/tree/master/official/resnet.\n",
    "\n",
    "To run this notebook, start the NGC TF container providing correct path to ImageNet validation data and a TF saved model:\n",
    "\n",
    "```bash\n",
    "nvidia-docker run -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",
    "This repository can then be cloned to `/workspace`:\n",
    "\n",
    "```bash\n",
    "git clone https://github.com/vinhngx/tftrt-examples\n",
    "```\n",
    "\n",
    "Then start jupyter notebook within the container with:\n",
    "\n",
    "```bash\n",
    "cd tftrt-examples\n",
    "jupyter notebook --ip 0.0.0.0 --port 8888  --allow-root\n",
    "```\n",
    "\n",
    "Connect to Jupyter notebook web interface from your local host http://localhost:8888. \n",
    "\n",
    "We first install some extra packages and external dependencies. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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": 10,
   "metadata": {},
   "outputs": [],
   "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 os\n",
    "os.environ['CUDA_VISIBLE_DEVICES']='1'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data\n",
    "We verify that the correct data folder has been mounted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "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",
    "\n",
    "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\n",
    "\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]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Saved model\n",
    "If not already downloaded, we will be downloading and working with a ResNet-50 v2 (fp32, Accuracy 76.47%) from https://github.com/tensorflow/models/tree/master/official/resnet. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The file '/saved_model/resnet_v2_fp32_savedmodel_NCHW.tar.gz' exists.\n"
     ]
    }
   ],
   "source": [
    "%%bash\n",
    "FILE=/saved_model/resnet_v2_fp32_savedmodel_NCHW.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_v2_fp32_savedmodel_NCHW.tar.gz \n",
    "fi\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./resnet_v2_fp32_savedmodel_NCHW/\n",
      "./resnet_v2_fp32_savedmodel_NCHW/1538687196/\n",
      "./resnet_v2_fp32_savedmodel_NCHW/1538687196/variables/\n",
      "./resnet_v2_fp32_savedmodel_NCHW/1538687196/variables/variables.index\n",
      "./resnet_v2_fp32_savedmodel_NCHW/1538687196/variables/variables.data-00000-of-00001\n",
      "./resnet_v2_fp32_savedmodel_NCHW/1538687196/saved_model.pb\n"
     ]
    }
   ],
   "source": [
    "!tar -xzvf /saved_model/resnet_v2_fp32_savedmodel_NCHW.tar.gz -C /saved_model "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We employ TF `saved_model_cli` tool to inspect the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "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/resnet_v2_fp32_savedmodel_NCHW/1538687196/"
   ]
  },
  {
   "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": 13,
   "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": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from preprocessing import inception_preprocessing, 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 (don't use 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": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def benchmark_saved_model(SAVED_MODEL_DIR, BATCH_SIZE):\n",
    "    with tf.Session(graph=tf.Graph()) as sess:\n",
    "        # Initialize all tfrecord paths\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=8))\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",
    "        num_hits = 0\n",
    "        num_predict = 0\n",
    "        print('Warming up for 10 batches...')\n",
    "        for _ in range (10):\n",
    "            image_data = sess.run(next_element)    \n",
    "            img = image_data[0]\n",
    "            label = image_data[1].squeeze()\n",
    "            output = sess.run(['softmax_tensor:0'], feed_dict={\"input_tensor:0\": img})\n",
    "            prediction = np.argmax(output[0], axis=1)\n",
    "            num_hits += np.sum(prediction == label)\n",
    "            num_predict += len(prediction)\n",
    "            \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(['softmax_tensor:0'], feed_dict={\"input_tensor:0\": 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": "markdown",
   "metadata": {},
   "source": [
    "## Benchmarking naitive Tensorflow model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Define some global variables\n",
    "BATCH_SIZE = 64\n",
    "SAVED_MODEL_DIR = \"/saved_model/resnet_v2_fp32_savedmodel_NCHW/1538687196/FP32\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Saver not created because there are no variables in the graph to restore\n",
      "INFO:tensorflow:The specified SavedModel has no variables; no checkpoints were restored.\n",
      "Warming up for 10 batches...\n",
      "Accuracy: 76.36%\n",
      "Inference speed: 1356.20 samples/s\n"
     ]
    }
   ],
   "source": [
    "benchmark_saved_model(SAVED_MODEL_DIR, BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Benchmarking TF-TRT FP32  inference engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Running against TensorRT version 5.1.2\n",
      "INFO:tensorflow:Saver not created because there are no variables in the graph to restore\n",
      "INFO:tensorflow:The specified SavedModel has no variables; no checkpoints were restored.\n",
      "INFO:tensorflow:Froze 0 variables.\n",
      "INFO:tensorflow:Converted 0 variables to const ops.\n",
      "INFO:tensorflow:No assets to save.\n",
      "INFO:tensorflow:No assets to write.\n",
      "INFO:tensorflow:SavedModel written to: /saved_model/resnet_v2_fp32_savedmodel_NCHW/1538687196/FP32FP32/saved_model.pb\n"
     ]
    }
   ],
   "source": [
    "FP32_SAVED_MODEL_DIR = SAVED_MODEL_DIR + \"FP32\"\n",
    "!rm -rf $FP32_SAVED_MODEL_DIR\n",
    "\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\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Saver not created because there are no variables in the graph to restore\n",
      "INFO:tensorflow:The specified SavedModel has no variables; no checkpoints were restored.\n",
      "Warming up for 10 batches...\n",
      "Accuracy: 76.36%\n",
      "Inference speed: 1365.15 samples/s\n"
     ]
    }
   ],
   "source": [
    "benchmark_saved_model(FP32_SAVED_MODEL_DIR, BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Benchmarking TF-TRT FP16 inference engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Running against TensorRT version 5.1.2\n",
      "INFO:tensorflow:Saver not created because there are no variables in the graph to restore\n",
      "INFO:tensorflow:The specified SavedModel has no variables; no checkpoints were restored.\n",
      "INFO:tensorflow:Froze 0 variables.\n",
      "INFO:tensorflow:Converted 0 variables to const ops.\n",
      "INFO:tensorflow:No assets to save.\n",
      "INFO:tensorflow:No assets to write.\n",
      "INFO:tensorflow:SavedModel written to: /saved_model/resnet_v2_fp32_savedmodel_NCHW/1538687196/FP32FP16/saved_model.pb\n"
     ]
    }
   ],
   "source": [
    "FP16_SAVED_MODEL_DIR = SAVED_MODEL_DIR + \"FP16\"\n",
    "!rm -rf $FP16_SAVED_MODEL_DIR\n",
    "\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\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Saver not created because there are no variables in the graph to restore\n",
      "INFO:tensorflow:The specified SavedModel has no variables; no checkpoints were restored.\n",
      "Warming up for 10 batches...\n",
      "Accuracy: 76.36%\n",
      "Inference speed: 1361.09 samples/s\n"
     ]
    }
   ],
   "source": [
    "benchmark_saved_model(FP16_SAVED_MODEL_DIR, BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating TFTRT INT8 inference 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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Saver not created because there are no variables in the graph to restore\n",
      "INFO:tensorflow:The specified SavedModel has no variables; no checkpoints were restored.\n",
      "INFO:tensorflow:Running against TensorRT version 5.1.2\n",
      "Calibrate model on calibration data...\n",
      "Calibration accuracy: 75.38%\n"
     ]
    }
   ],
   "source": [
    "tf.reset_default_graph()\n",
    "with tf.Session(graph=tf.Graph()) as sess:\n",
    "    tf.saved_model.loader.load(\n",
    "        sess, [tf.saved_model.tag_constants.SERVING], SAVED_MODEL_DIR)\n",
    "    fp32_graph_def = sess.graph_def\n",
    "len(fp32_graph_def.SerializeToString())    \n",
    "\n",
    "#Now we create the TFTRT INT8 calibration graph\n",
    "trt_int8_calib_graph = trt.create_inference_graph(\n",
    "        input_graph_def=fp32_graph_def,\n",
    "        outputs=['softmax_tensor:0'],\n",
    "        max_batch_size=BATCH_SIZE,\n",
    "        max_workspace_size_bytes=1<<32,\n",
    "        precision_mode='INT8')\n",
    "\n",
    "#Then calibrate it with 2 batchs of examples\n",
    "N_runs=100\n",
    "with tf.Session(graph=tf.Graph()) as sess:\n",
    "    \n",
    "    output_node = tf.import_graph_def(\n",
    "        trt_int8_calib_graph,\n",
    "        return_elements=['softmax_tensor'],\n",
    "        name='trt')\n",
    "    \n",
    "    # Initialize all tfrecord paths\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=8))\n",
    "    iterator = dataset.make_one_shot_iterator()\n",
    "    next_element = iterator.get_next()\n",
    "\n",
    "\n",
    "    #print([n.name for n in tf.get_default_graph().as_graph_def().node])\n",
    "    \n",
    "    print('Calibrate model on calibration data...')\n",
    "    num_hits = 0\n",
    "    num_predict = 0\n",
    "    for _ in range(N_runs):\n",
    "            image_data = sess.run(next_element)    \n",
    "            img = image_data[0]\n",
    "            label = image_data[1].squeeze()\n",
    "            output = sess.run(output_node[0].outputs[0], feed_dict={\"trt/input_tensor:0\": img})\n",
    "            prediction = np.argmax(output, axis=1)\n",
    "            num_hits += np.sum(prediction == label)\n",
    "            num_predict += len(prediction)\n",
    "    print('Calibration accuracy: %.2f%%'%(100*num_hits/num_predict)) \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2\n",
    "\n",
    "Now we convert the INT8 calibration graph to the final TF-TRT INT8 inference engine, and benchmark its performance. We will also be saving this engine to a *saved model*, ready to be served elsewhere."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "INT8_SAVED_MODEL_DIR = SAVED_MODEL_DIR + \"INT8\"\n",
    "!rm -rf $INT8_SAVED_MODEL_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "ename": "FailedPreconditionError",
     "evalue": "Need to run graph with calibration data first!",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFailedPreconditionError\u001b[0m                   Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-38-774a4b39c75b>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m#Create Int8 inference model from the calibration graph and write to a saved session\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mtrt_int8_calibrated_graph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtrt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcalib_graph_to_infer_graph\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrt_int8_calib_graph\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m output_node = tf.import_graph_def(\n\u001b[1;32m      4\u001b[0m         \u001b[0mtrt_int8_calibrated_graph\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m         \u001b[0mreturn_elements\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'softmax_tensor:0'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.5/dist-packages/tensorflow/contrib/tensorrt/python/trt_convert.py\u001b[0m in \u001b[0;36mcalib_graph_to_infer_graph\u001b[0;34m(calibration_graph_def, is_dynamic_op)\u001b[0m\n\u001b[1;32m    415\u001b[0m     \u001b[0;31m# pylint: disable=protected-access\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    416\u001b[0m     raise _impl._make_specific_exception(None, None, \";\".join(msg[1:]),\n\u001b[0;32m--> 417\u001b[0;31m                                          int(msg[0]))\n\u001b[0m\u001b[1;32m    418\u001b[0m     \u001b[0;31m# pylint: enable=protected-access\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    419\u001b[0m   \u001b[0moutput_graph_def\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgraph_pb2\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mGraphDef\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mFailedPreconditionError\u001b[0m: Need to run graph with calibration data first!"
     ]
    }
   ],
   "source": [
    "#Create Int8 inference model from the calibration graph and write to a saved session\n",
    "trt_int8_calibrated_graph=trt.calib_graph_to_infer_graph(trt_int8_calib_graph)\n",
    "output_node = tf.import_graph_def(\n",
    "        trt_int8_calibrated_graph,\n",
    "        return_elements=['softmax_tensor:0'],\n",
    "        name='')\n",
    "\n",
    "with tf.Session(graph=tf.Graph()) as sess:\n",
    "    # Initialize all tfrecord paths\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=8))\n",
    "    iterator = dataset.make_one_shot_iterator()\n",
    "    next_element = iterator.get_next()\n",
    "\n",
    "    output_node = tf.import_graph_def(\n",
    "        trt_int8_calibrated_graph,\n",
    "        return_elements=['softmax_tensor:0'],\n",
    "        name='')\n",
    "\n",
    "    print('Warming up for 10 batches...')\n",
    "    for _ in range (10):\n",
    "        image_data = sess.run(next_element)    \n",
    "        img = image_data[0]\n",
    "        output = sess.run(['softmax_tensor:0'], feed_dict={\"input_tensor:0\": img})\n",
    "        \n",
    "    print('Benchmarking TF-TRT INT8 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(['softmax_tensor:0'], feed_dict={\"input_tensor:0\": 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('INT8 TF-TRT Accuracy: %.2f%%'%(100*num_hits/num_predict))\n",
    "    print('INT8 TF-TRT Inference speed: %.2f samples/s'%(num_predict/(time.time()-start_time)))\n",
    "\n",
    "    #Save model for serving\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(\"import/input:0\")},\n",
    "        outputs={\"softmax\":output_node[0].outputs[0]},\n",
    "        legacy_init_op=None\n",
    "     )\n",
    "\n",
    "      "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Benchmarking INT8 saved model\n",
    "\n",
    "Finally we reload and verify the performance of the INT8 saved model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#benchmark int8 saved model\n",
    "with tf.Session(graph=tf.Graph()) as sess:\n",
    "    # Initialize all tfrecord paths\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=8))\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_INT8_MODEL_DIR)\n",
    "    \n",
    "    print('Warming up for 10 batches...')\n",
    "    for _ in range (10):\n",
    "        image_data = sess.run(next_element)    \n",
    "        img = image_data[0]\n",
    "        output = sess.run(['import/resnet_v1_50/predictions/Softmax:0'], feed_dict={\"import/input:0\": img})\n",
    "    \n",
    "    print('Benchmarking TF-TRT INT8 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]\n",
    "            output = sess.run(['import/resnet_v1_50/predictions/Softmax:0'], feed_dict={\"import/input:0\": 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('INT8 TF-TRT Accuracy: %.2f%%'%(100*num_hits/num_predict))\n",
    "    print('INT8 TF-TRT Inference speed: %.2f samples/s'%(num_predict/(time.time()-start_time)))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Benchmarking with synthetic data\n",
    "\n",
    "While benchmarking with real datasets, there are data reading and pre-processing procedures involved. As a result, the GPU is not fully loaded all the time. In this section, we test with synthetic data to test the throughput limit of the GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_ITER = 100\n",
    "dummy_input = np.random.random_sample((BATCH_SIZE,224,224,3))\n",
    "\n",
    "fp32_graph = tf.Graph()\n",
    "with fp32_graph.as_default():\n",
    "    inc=tf.constant(dummy_input, dtype=tf.float32)\n",
    "    dataset=tf.data.Dataset.from_tensors(inc)\n",
    "    dataset=dataset.repeat()\n",
    "    iterator=dataset.make_one_shot_iterator()\n",
    "    next_element=iterator.get_next()\n",
    "    out = tf.import_graph_def(\n",
    "      graph_def=fp32_graph_def,\n",
    "      input_map={\"import/input\":next_element},\n",
    "      return_elements=[ \"import/resnet_v1_50/predictions/Softmax\"]\n",
    "    )\n",
    "\n",
    "with tf.Session(graph=fp32_graph) as sess:\n",
    "    print('Warming up for 10 batches...')\n",
    "    for _ in range(10):\n",
    "        sess.run(out)\n",
    "    \n",
    "    print('Benchmarking...')\n",
    "    start_time = time.time()\n",
    "    for _ in range(NUM_ITER):\n",
    "        sess.run(out)\n",
    "    print('Naitive FP32 Inference speed: %.2f samples/s'%(NUM_ITER*BATCH_SIZE/(time.time()-start_time)))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fp16_graph = tf.Graph()\n",
    "with fp16_graph.as_default():\n",
    "    inc=tf.constant(dummy_input, dtype=tf.float32)\n",
    "    dataset=tf.data.Dataset.from_tensors(inc)\n",
    "    dataset=dataset.repeat()\n",
    "    iterator=dataset.make_one_shot_iterator()\n",
    "    next_element=iterator.get_next()\n",
    "    out = tf.import_graph_def(\n",
    "      graph_def=trt_fp16_Graph,\n",
    "      input_map={\"import/input\":next_element},\n",
    "      return_elements=[ \"import/resnet_v1_50/predictions/Softmax\"]\n",
    "    )\n",
    "    out = out[0].outputs[0]\n",
    "\n",
    "with tf.Session(graph=fp16_graph) as sess:\n",
    "    print('Warming up for 10 batches...')\n",
    "    for _ in range(10):\n",
    "        sess.run(out)\n",
    "    \n",
    "    print('Benchmarking...')\n",
    "    start_time = time.time()\n",
    "    for _ in range(NUM_ITER):\n",
    "        sess.run(out)\n",
    "    print('Naitive FP32 Inference speed: %.2f samples/s'%(NUM_ITER*BATCH_SIZE/(time.time()-start_time)))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "int8_graph = tf.Graph()\n",
    "with int8_graph.as_default():\n",
    "    inc=tf.constant(dummy_input, dtype=tf.float32)\n",
    "    dataset=tf.data.Dataset.from_tensors(inc)\n",
    "    dataset=dataset.repeat()\n",
    "    iterator=dataset.make_one_shot_iterator()\n",
    "    next_element=iterator.get_next()\n",
    "    out = tf.import_graph_def(\n",
    "      graph_def=trt_int8_calibrated_graph,\n",
    "      input_map={\"import/input\":next_element},\n",
    "      return_elements=[ \"import/resnet_v1_50/predictions/Softmax\"],\n",
    "      name=''        \n",
    "    )\n",
    "        \n",
    "with tf.Session(graph=int8_graph) as sess:\n",
    "\n",
    "    print('Warming up for 10 batches...')\n",
    "    for _ in range(10):\n",
    "        sess.run(out)\n",
    "    \n",
    "    print('Benchmarking...')\n",
    "    start_time = time.time()\n",
    "    for _ in range(NUM_ITER):\n",
    "        sess.run(out)\n",
    "    print('Naitive FP32 Inference speed: %.2f samples/s'%(NUM_ITER*BATCH_SIZE/(time.time()-start_time)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}