Step1_Object_detection_Colab_TensorRT.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Step1_Object_detection_Colab_TensorRT.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Tony607/tf_jetson_nano/blob/master/Step1_Object_detection_Colab_TensorRT.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "metadata": {
        "id": "QQK7KlhRBfnf",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# Step 1 : Convert TensorFlow object detection model into TensorRT model\n",
        "**For more detail, checkout [How to run TensorFlow Object Detection model on Jetson Nano](https://www.dlology.com/blog/how-to-run-tensorflow-object-detection-model-on-jetson-nano/) | DLology Blog**"
      ]
    },
    {
      "metadata": {
        "id": "T_-qqpBAM3LW",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Pre-trained TensorFlow object detection models"
      ]
    },
    {
      "metadata": {
        "id": "ucxso6LcMows",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "from collections import namedtuple\n",
        "DetectionModel = namedtuple('DetectionModel', ['name', 'url', 'extract_dir'])\n",
        "\n",
        "MODELS = {\n",
        "    'ssd_mobilenet_v1_coco': DetectionModel(\n",
        "        'ssd_mobilenet_v1_coco',\n",
        "        'http://download.tensorflow.org/models/object_detection/ssd_mobilenet_v1_coco_2018_01_28.tar.gz',\n",
        "        'ssd_mobilenet_v1_coco_2018_01_28',\n",
        "    ),\n",
        "    'ssd_mobilenet_v2_coco': DetectionModel(\n",
        "        'ssd_mobilenet_v2_coco',\n",
        "        'http://download.tensorflow.org/models/object_detection/ssd_mobilenet_v2_coco_2018_03_29.tar.gz',\n",
        "        'ssd_mobilenet_v2_coco_2018_03_29',\n",
        "    ),\n",
        "    'ssd_inception_v2_coco': DetectionModel(\n",
        "        'ssd_inception_v2_coco',\n",
        "        'http://download.tensorflow.org/models/object_detection/ssd_inception_v2_coco_2018_01_28.tar.gz',\n",
        "        'ssd_inception_v2_coco_2018_01_28',\n",
        "    ),\n",
        "    'ssd_resnet_50_fpn_coco': DetectionModel(\n",
        "        'ssd_resnet_50_fpn_coco',\n",
        "        'http://download.tensorflow.org/models/object_detection/ssd_resnet50_v1_fpn_shared_box_predictor_640x640_coco14_sync_2018_07_03.tar.gz',\n",
        "        'ssd_resnet50_v1_fpn_shared_box_predictor_640x640_coco14_sync_2018_07_03',\n",
        "    ),\n",
        "    'faster_rcnn_resnet50_coco': DetectionModel(\n",
        "        'faster_rcnn_resnet50_coco',\n",
        "        'http://download.tensorflow.org/models/object_detection/faster_rcnn_resnet50_coco_2018_01_28.tar.gz',\n",
        "        'faster_rcnn_resnet50_coco_2018_01_28',\n",
        "    ),\n",
        "    'faster_rcnn_nas': DetectionModel(\n",
        "        'faster_rcnn_nas',\n",
        "        'http://download.tensorflow.org/models/object_detection/faster_rcnn_nas_coco_2018_01_28.tar.gz',\n",
        "        'faster_rcnn_nas_coco_2018_01_28',\n",
        "    ),\n",
        "    'mask_rcnn_resnet50_atrous_coco': DetectionModel(\n",
        "        'mask_rcnn_resnet50_atrous_coco',\n",
        "        'http://download.tensorflow.org/models/object_detection/mask_rcnn_resnet50_atrous_coco_2018_01_28.tar.gz',\n",
        "        'mask_rcnn_resnet50_atrous_coco_2018_01_28',\n",
        "    )\n",
        "}\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "UyEDVcwfNDO-",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Select your model"
      ]
    },
    {
      "metadata": {
        "id": "1sENfWFOIiTr",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "MODEL = 'ssd_mobilenet_v1_coco'"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "MI6QhZndMZFV",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Install required packages"
      ]
    },
    {
      "metadata": {
        "id": "GrcV6tNTDT8u",
        "colab_type": "code",
        "outputId": "f78f6b33-7004-4b0d-d1c1-7f61e79f2a63",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 290
        }
      },
      "cell_type": "code",
      "source": [
        "%cd /content\n",
        "!git clone --quiet https://github.com/tensorflow/models.git\n",
        "\n",
        "!apt-get install -qq protobuf-compiler python-pil python-lxml python-tk\n",
        "\n",
        "!pip install -q Cython contextlib2 pillow lxml matplotlib\n",
        "\n",
        "!pip install -q pycocotools\n",
        "\n",
        "%cd /content/models/research\n",
        "!protoc object_detection/protos/*.proto --python_out=.\n",
        "\n",
        "import os\n",
        "import sys\n",
        "os.environ['PYTHONPATH'] += ':/content/models/research/:/content/models/research/slim/'\n",
        "sys.path.append(\"/content/models/research/slim/\")\n",
        "\n",
        "!python object_detection/builders/model_builder_test.py"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "/content\n",
            "fatal: destination path 'models' already exists and is not an empty directory.\n",
            "/content/models/research\n",
            "\n",
            "WARNING: The TensorFlow contrib module will not be included in TensorFlow 2.0.\n",
            "For more information, please see:\n",
            "  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n",
            "  * https://github.com/tensorflow/addons\n",
            "If you depend on functionality not listed there, please file an issue.\n",
            "\n",
            "............s...\n",
            "----------------------------------------------------------------------\n",
            "Ran 16 tests in 0.062s\n",
            "\n",
            "OK (skipped=1)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "jZVU5IGWLq87",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "import tensorflow.contrib.tensorrt as trt"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "l9UHOzwHNsBd",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "`graph_utils.py`"
      ]
    },
    {
      "metadata": {
        "id": "axmgTJ8BINlN",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "\n",
        "def make_const6(const6_name='const6'):\n",
        "    graph = tf.Graph()\n",
        "    with graph.as_default():\n",
        "        tf_6 = tf.constant(dtype=tf.float32, value=6.0, name=const6_name)\n",
        "    return graph.as_graph_def()\n",
        "\n",
        "\n",
        "def make_relu6(output_name, input_name, const6_name='const6'):\n",
        "    graph = tf.Graph()\n",
        "    with graph.as_default():\n",
        "        tf_x = tf.placeholder(tf.float32, [10, 10], name=input_name)\n",
        "        tf_6 = tf.constant(dtype=tf.float32, value=6.0, name=const6_name)\n",
        "        with tf.name_scope(output_name):\n",
        "            tf_y1 = tf.nn.relu(tf_x, name='relu1')\n",
        "            tf_y2 = tf.nn.relu(tf.subtract(tf_x, tf_6, name='sub1'), name='relu2')\n",
        "\n",
        "            #tf_y = tf.nn.relu(tf.subtract(tf_6, tf.nn.relu(tf_x, name='relu1'), name='sub'), name='relu2')\n",
        "        #tf_y = tf.subtract(tf_6, tf_y, name=output_name)\n",
        "        tf_y = tf.subtract(tf_y1, tf_y2, name=output_name)\n",
        "        \n",
        "    graph_def = graph.as_graph_def()\n",
        "    graph_def.node[-1].name = output_name\n",
        "\n",
        "    # remove unused nodes\n",
        "    for node in graph_def.node:\n",
        "        if node.name == input_name:\n",
        "            graph_def.node.remove(node)\n",
        "    for node in graph_def.node:\n",
        "        if node.name == const6_name:\n",
        "            graph_def.node.remove(node)\n",
        "    for node in graph_def.node:\n",
        "        if node.op == '_Neg':\n",
        "            node.op = 'Neg'\n",
        "            \n",
        "    return graph_def\n",
        "\n",
        "\n",
        "def convert_relu6(graph_def, const6_name='const6'):\n",
        "    # add constant 6\n",
        "    has_const6 = False\n",
        "    for node in graph_def.node:\n",
        "        if node.name == const6_name:\n",
        "            has_const6 = True\n",
        "    if not has_const6:\n",
        "        const6_graph_def = make_const6(const6_name=const6_name)\n",
        "        graph_def.node.extend(const6_graph_def.node)\n",
        "        \n",
        "    for node in graph_def.node:\n",
        "        if node.op == 'Relu6':\n",
        "            input_name = node.input[0]\n",
        "            output_name = node.name\n",
        "            relu6_graph_def = make_relu6(output_name, input_name, const6_name=const6_name)\n",
        "            graph_def.node.remove(node)\n",
        "            graph_def.node.extend(relu6_graph_def.node)\n",
        "            \n",
        "    return graph_def\n",
        "\n",
        "\n",
        "def remove_node(graph_def, node):\n",
        "    for n in graph_def.node:\n",
        "        if node.name in n.input:\n",
        "            n.input.remove(node.name)\n",
        "        ctrl_name = '^' + node.name\n",
        "        if ctrl_name in n.input:\n",
        "            n.input.remove(ctrl_name)\n",
        "    graph_def.node.remove(node)\n",
        "\n",
        "\n",
        "def remove_op(graph_def, op_name):\n",
        "    matches = [node for node in graph_def.node if node.op == op_name]\n",
        "    for match in matches:\n",
        "        remove_node(graph_def, match)\n",
        "\n",
        "\n",
        "def f_force_nms_cpu(frozen_graph):\n",
        "    for node in frozen_graph.node:\n",
        "        if 'NonMaxSuppression' in node.name:\n",
        "            node.device = '/device:CPU:0'\n",
        "    return frozen_graph\n",
        "\n",
        "\n",
        "def f_replace_relu6(frozen_graph):\n",
        "    return convert_relu6(frozen_graph)\n",
        "\n",
        "\n",
        "def f_remove_assert(frozen_graph):\n",
        "    remove_op(frozen_graph, 'Assert')\n",
        "    return frozen_graph\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "dm7-PKVDN1Ll",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "`detection.py`"
      ]
    },
    {
      "metadata": {
        "id": "ihIdQpBFIKE4",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "from object_detection.protos import pipeline_pb2\n",
        "from object_detection import exporter\n",
        "\n",
        "import os\n",
        "import subprocess\n",
        "\n",
        "from google.protobuf import text_format\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "\n",
        "INPUT_NAME='image_tensor'\n",
        "BOXES_NAME='detection_boxes'\n",
        "CLASSES_NAME='detection_classes'\n",
        "SCORES_NAME='detection_scores'\n",
        "MASKS_NAME='detection_masks'\n",
        "NUM_DETECTIONS_NAME='num_detections'\n",
        "FROZEN_GRAPH_NAME='frozen_inference_graph.pb'\n",
        "PIPELINE_CONFIG_NAME='pipeline.config'\n",
        "CHECKPOINT_PREFIX='model.ckpt'\n",
        "\n",
        "\n",
        "\n",
        "def get_input_names(model):\n",
        "    return [INPUT_NAME]\n",
        "\n",
        "\n",
        "def get_output_names(model):\n",
        "    output_names = [BOXES_NAME, CLASSES_NAME, SCORES_NAME, NUM_DETECTIONS_NAME]\n",
        "    if model == 'mask_rcnn_resnet50_atrous_coco':\n",
        "        output_names.append(MASKS_NAME)\n",
        "    return output_names\n",
        "\n",
        "\n",
        "def download_detection_model(model, output_dir='.'):\n",
        "    \"\"\"Downloads a pre-trained object detection model\"\"\"\n",
        "    global MODELS\n",
        "\n",
        "    model_name = model\n",
        "\n",
        "    model = MODELS[model_name]\n",
        "    subprocess.call(['mkdir', '-p', output_dir])\n",
        "    tar_file = os.path.join(output_dir, os.path.basename(model.url))\n",
        "\n",
        "    config_path = os.path.join(output_dir, model.extract_dir, PIPELINE_CONFIG_NAME)\n",
        "    checkpoint_path = os.path.join(output_dir, model.extract_dir, CHECKPOINT_PREFIX)\n",
        "\n",
        "    if not os.path.exists(os.path.join(output_dir, model.extract_dir)):\n",
        "        subprocess.call(['wget', model.url, '-O', tar_file])\n",
        "        subprocess.call(['tar', '-xzf', tar_file, '-C', output_dir])\n",
        "\n",
        "        # hack fix to handle mobilenet_v2 config bug\n",
        "        subprocess.call(['sed', '-i', '/batch_norm_trainable/d', config_path])\n",
        "\n",
        "    return config_path, checkpoint_path\n",
        "\n",
        "\n",
        "def build_detection_graph(config, checkpoint,\n",
        "        batch_size=1,\n",
        "        score_threshold=None,\n",
        "        iou_threshold=None,\n",
        "        force_nms_cpu=True,\n",
        "        replace_relu6=True,\n",
        "        remove_assert=True,\n",
        "        input_shape=None,\n",
        "        output_dir='.generated_model'):\n",
        "    \"\"\"Builds a frozen graph for a pre-trained object detection model\"\"\"\n",
        "    \n",
        "    config_path = config\n",
        "    checkpoint_path = checkpoint\n",
        "\n",
        "    # parse config from file\n",
        "    config = pipeline_pb2.TrainEvalPipelineConfig()\n",
        "    with open(config_path, 'r') as f:\n",
        "        text_format.Merge(f.read(), config, allow_unknown_extension=True)\n",
        "\n",
        "    # override some config parameters\n",
        "    if config.model.HasField('ssd'):\n",
        "        config.model.ssd.feature_extractor.override_base_feature_extractor_hyperparams = True\n",
        "        if score_threshold is not None:\n",
        "            config.model.ssd.post_processing.batch_non_max_suppression.score_threshold = score_threshold\n",
        "        if iou_threshold is not None:\n",
        "            config.model.ssd.post_processing.batch_non_max_suppression.iou_threshold = iou_threshold\n",
        "        if input_shape is not None:\n",
        "            config.model.ssd.image_resizer.fixed_shape_resizer.height = input_shape[0]\n",
        "            config.model.ssd.image_resizer.fixed_shape_resizer.width = input_shape[1]\n",
        "    elif config.model.HasField('faster_rcnn'):\n",
        "        if score_threshold is not None:\n",
        "            config.model.faster_rcnn.second_stage_post_processing.score_threshold = score_threshold\n",
        "        if input_shape is not None:\n",
        "            config.model.faster_rcnn.image_resizer.fixed_shape_resizer.height = input_shape[0]\n",
        "            config.model.faster_rcnn.image_resizer.fixed_shape_resizer.width = input_shape[1]\n",
        "\n",
        "    if os.path.isdir(output_dir):\n",
        "        subprocess.call(['rm', '-rf', output_dir])\n",
        "\n",
        "    tf_config = tf.ConfigProto()\n",
        "    tf_config.gpu_options.allow_growth = True\n",
        "\n",
        "    # export inference graph to file (initial)\n",
        "    with tf.Session(config=tf_config) as tf_sess:\n",
        "        with tf.Graph().as_default() as tf_graph:\n",
        "            exporter.export_inference_graph(\n",
        "                'image_tensor', \n",
        "                config, \n",
        "                checkpoint_path, \n",
        "                output_dir, \n",
        "                input_shape=[batch_size, None, None, 3]\n",
        "            )\n",
        "\n",
        "    # read frozen graph from file\n",
        "    frozen_graph = tf.GraphDef()\n",
        "    with open(os.path.join(output_dir, FROZEN_GRAPH_NAME), 'rb') as f:\n",
        "        frozen_graph.ParseFromString(f.read())\n",
        "\n",
        "    # apply graph modifications\n",
        "    if force_nms_cpu:\n",
        "        frozen_graph = f_force_nms_cpu(frozen_graph)\n",
        "    if replace_relu6:\n",
        "        frozen_graph = f_replace_relu6(frozen_graph)\n",
        "    if remove_assert:\n",
        "        frozen_graph = f_remove_assert(frozen_graph)\n",
        "\n",
        "    # get input names\n",
        "    # TODO: handle mask_rcnn \n",
        "    input_names = [INPUT_NAME]\n",
        "    output_names = [BOXES_NAME, CLASSES_NAME, SCORES_NAME, NUM_DETECTIONS_NAME]\n",
        "\n",
        "    # remove temporary directory\n",
        "    subprocess.call(['rm', '-rf', output_dir])\n",
        "\n",
        "    return frozen_graph, input_names, output_names\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "KLEiAFOAMNJ8",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Download pre-trained TensorFlow Object detection model"
      ]
    },
    {
      "metadata": {
        "id": "f7WONX1MIp0k",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "config_path, checkpoint_path = download_detection_model(MODEL, 'data')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "fpj3sFQVMUFA",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "For improved performance, increase the non-max suppression score threshold in the downloaded config file from 1e-8 to something greater, like 0.1."
      ]
    },
    {
      "metadata": {
        "id": "vdl8C-V8IrIA",
        "colab_type": "code",
        "outputId": "976b27c6-e8ee-4114-de52-94c09bf593b1",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 308
        }
      },
      "cell_type": "code",
      "source": [
        "frozen_graph, input_names, output_names = build_detection_graph(\n",
        "    config=config_path,\n",
        "    checkpoint=checkpoint_path,\n",
        "    score_threshold=0.3,\n",
        "    iou_threshold=0.5,\n",
        "    batch_size=1\n",
        ")"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:depth of additional conv before box predictor: 0\n",
            "INFO:tensorflow:depth of additional conv before box predictor: 0\n",
            "INFO:tensorflow:depth of additional conv before box predictor: 0\n",
            "INFO:tensorflow:depth of additional conv before box predictor: 0\n",
            "INFO:tensorflow:depth of additional conv before box predictor: 0\n",
            "INFO:tensorflow:depth of additional conv before box predictor: 0\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "18 ops no flops stats due to incomplete shapes.\n",
            "18 ops no flops stats due to incomplete shapes.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Restoring parameters from data/ssd_mobilenet_v1_coco_2018_01_28/model.ckpt\n",
            "INFO:tensorflow:Restoring parameters from data/ssd_mobilenet_v1_coco_2018_01_28/model.ckpt\n",
            "INFO:tensorflow:Froze 199 variables.\n",
            "INFO:tensorflow:Converted 199 variables to const ops.\n",
            "INFO:tensorflow:No assets to save.\n",
            "INFO:tensorflow:No assets to write.\n",
            "INFO:tensorflow:SavedModel written to: .generated_model/saved_model/saved_model.pb\n",
            "INFO:tensorflow:Writing pipeline config file to .generated_model/pipeline.config\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "eWyu5YhLIrKe",
        "colab_type": "code",
        "outputId": "65331419-4fbe-4234-c6ee-4d781d7926b9",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "cell_type": "code",
      "source": [
        "print(output_names)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "['detection_boxes', 'detection_classes', 'detection_scores', 'num_detections']\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "yBbfUCvEODTo",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Optimize the model with TensorRT"
      ]
    },
    {
      "metadata": {
        "id": "uKC8ZWZEIwaS",
        "colab_type": "code",
        "outputId": "ba26eff1-0851-4306-e247-056f406fbd1f",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "cell_type": "code",
      "source": [
        "trt_graph = trt.create_inference_graph(\n",
        "    input_graph_def=frozen_graph,\n",
        "    outputs=output_names,\n",
        "    max_batch_size=1,\n",
        "    max_workspace_size_bytes=1 << 25,\n",
        "    precision_mode='FP16',\n",
        "    minimum_segment_size=50\n",
        ")"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Running against TensorRT version 0.0.0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "5vUsC1mCIwdz",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "with open('./data/trt_graph.pb', 'wb') as f:\n",
        "    f.write(trt_graph.SerializeToString())"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "KOgeJ-lyIwgU",
        "colab_type": "code",
        "outputId": "42bb3f5d-41e9-492c-fc07-052ddc4e0749",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 126
        }
      },
      "cell_type": "code",
      "source": [
        "!ls data -alh"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "total 101M\n",
            "drwxr-xr-x  3 root   root 4.0K Apr 20 05:19 .\n",
            "drwxr-xr-x 72 root   root 4.0K Apr 20 07:20 ..\n",
            "drwxr-xr-x  3 345018 5000 4.0K Apr 20 05:05 ssd_mobilenet_v1_coco_2018_01_28\n",
            "-rw-r--r--  1 root   root  73M Feb 10  2018 ssd_mobilenet_v1_coco_2018_01_28.tar.gz\n",
            "-rw-r--r--  1 root   root  28M Apr 20 07:20 trt_graph.pb\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "jWDYMW6ZAjbw",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Benchmark TensorFlow prediction speed"
      ]
    },
    {
      "metadata": {
        "id": "RwUKKaM2_D95",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "input_names = ['image_tensor']\n",
        "output_names = ['detection_boxes', 'detection_classes', 'detection_scores', 'num_detections']\n",
        "# Create session and load graph\n",
        "tf_config = tf.ConfigProto()\n",
        "tf_config.gpu_options.allow_growth = True\n",
        "tf_sess = tf.Session(config=tf_config)\n",
        "tf.import_graph_def(frozen_graph, name='')\n",
        "\n",
        "tf_input = tf_sess.graph.get_tensor_by_name(input_names[0] + ':0')\n",
        "tf_scores = tf_sess.graph.get_tensor_by_name('detection_scores:0')\n",
        "tf_boxes = tf_sess.graph.get_tensor_by_name('detection_boxes:0')\n",
        "tf_classes = tf_sess.graph.get_tensor_by_name('detection_classes:0')\n",
        "tf_num_detections = tf_sess.graph.get_tensor_by_name('num_detections:0')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "ie0zwepg_MJA",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "image = np.random.random((300,300,3))\n",
        "scores, boxes, classes, num_detections = tf_sess.run([tf_scores, tf_boxes, tf_classes, tf_num_detections], feed_dict={\n",
        "    tf_input: image[None, ...]\n",
        "})\n",
        "\n",
        "boxes = boxes[0] # index by 0 to remove batch dimension\n",
        "scores = scores[0]\n",
        "classes = classes[0]\n",
        "num_detections = num_detections[0]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "1e6NHFOWAREp",
        "colab_type": "code",
        "outputId": "f4ebf3b0-6e39-4f8d-b497-57ae82af0186",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "cell_type": "code",
      "source": [
        "import time\n",
        "times = []\n",
        "for i in range(20):\n",
        "    start_time = time.time()\n",
        "    scores, boxes, classes, num_detections = tf_sess.run([tf_scores, tf_boxes, tf_classes, tf_num_detections], feed_dict={\n",
        "        tf_input: image[None, ...]\n",
        "    })\n",
        "\n",
        "    delta = (time.time() - start_time)\n",
        "    times.append(delta)\n",
        "mean_delta = np.array(times).mean()\n",
        "fps = 1/mean_delta\n",
        "print('average(sec):{},fps:{}'.format(mean_delta,fps))\n"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "average(sec):0.010851287841796875,fps:92.15496027560994\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "ySxWp6zQBFL7",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "tf_sess.close()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "jPVLp1FKvnSF",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Download the tensorRT graph `.pb` file from colab to your local machine."
      ]
    },
    {
      "metadata": {
        "id": "ZrHyjN_Cvk4Z",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "from google.colab import files\n",
        "\n",
        "files.download('./data/trt_graph.pb')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "_Hg6qDGIwmQn",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "**Next step**: transfer the `trt_graph.pb` to your Jetson Nano, load it up and make predictions.\n",
        "\n",
        "\n",
        "`Step2_jetson-object-detection-predict.ipynb`"
      ]
    },
    {
      "metadata": {
        "id": "UkT7r-rnNj5j",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}