TensorFlow->TensorRT Image Classification

This contains examples, scripts and code related to image classification using TensorFlow models (from here) converted to TensorRT. Converting TensorFlow models to TensorRT offers significant performance gains on the Jetson TX2 as seen below.

Models
Setup
Download models and create frozen graphs
Convert frozen graph to TensorRT engine
Execute TensorRT engine
Benchmark all models

Models

The table below shows various details related to pretrained models ported from the TensorFlow slim model zoo.

_Model	_{Input Size}	_{TensorRT (TX2 / Half)}	_{TensorRT (TX2 / Float)}	_{TensorFlow (TX2 / Float)}	_{Input Name}	_{Output Name}	_{Preprocessing Fn.}
_{inception_v1}	_224x224	_7.98ms	_12.8ms	_27.6ms	_input	_{InceptionV1/Logits/SpatialSqueeze}	_inception
_{inception_v3}	_299x299	_26.3ms	_46.1ms	_98.4ms	_input	_{InceptionV3/Logits/SpatialSqueeze}	_inception
_{inception_v4}	_299x299	_52.1ms	_88.2ms	_176ms	_input	_{InceptionV4/Logits/Logits/BiasAdd}	_inception
_{inception_resnet_v2}	_299x299	_53.0ms	_98.7ms	_168ms	_input	_{InceptionResnetV2/Logits/Logits/BiasAdd}	_inception
_{resnet_v1_50}	_224x224	_15.7ms	_27.1ms	_63.9ms	_input	_{resnet_v1_50/SpatialSqueeze}	_vgg
_{resnet_v1_101}	_224x224	_29.9ms	_51.8ms	_107ms	_input	_{resnet_v1_101/SpatialSqueeze}	_vgg
_{resnet_v1_152}	_224x224	_42.6ms	_78.2ms	_157ms	_input	_{resnet_v1_152/SpatialSqueeze}	_vgg
_{resnet_v2_50}	_299x299	_27.5ms	_44.4ms	_92.2ms	_input	_{resnet_v2_50/SpatialSqueeze}	_inception
_{resnet_v2_101}	_299x299	_49.2ms	_83.1ms	_160ms	_input	_{resnet_v2_101/SpatialSqueeze}	_inception
_{resnet_v2_152}	_299x299	_74.6ms	_124ms	_230ms	_input	_{resnet_v2_152/SpatialSqueeze}	_inception
_{mobilenet_v1_0p25_128}	_128x128	_2.67ms	_2.65ms	_15.7ms	_input	_{MobilenetV1/Logits/SpatialSqueeze}	_inception
_{mobilenet_v1_0p5_160}	_160x160	_3.95ms	_4.00ms	_16.9ms	_input	_{MobilenetV1/Logits/SpatialSqueeze}	_inception
_{mobilenet_v1_1p0_224}	_224x224	_12.9ms	_12.9ms	_24.4ms	_input	_{MobilenetV1/Logits/SpatialSqueeze}	_inception
_{vgg_16}	_224x224	_38.2ms	_79.2ms	_171ms	_input	_{vgg_16/fc8/BiasAdd}	_vgg

The times recorded include data transfer to GPU, network execution, and data transfer back from GPU. Time does not include preprocessing. See scripts/test_tf.py, scripts/test_trt.py, and src/test/test_trt.cu for implementation details.

Setup

Flash the Jetson TX2 using JetPack 3.2. Be sure to install
- CUDA 9.0
- OpenCV4Tegra
- cuDNN
- TensorRT 3.0
Install pip on Jetson TX2.
```
sudo apt-get install python-pip
```
Install TensorFlow on Jetson TX2.
1. Download the TensorFlow 1.5.0 pip wheel from here. This build of TensorFlow is provided as a convenience for the purposes of this project.
2. Install TensorFlow using pip
```
  sudo pip install tensorflow-1.5.0rc0-cp27-cp27mu-linux_aarch64.whl
```
Install uff exporter on Jetson TX2.
1. Download TensorRT 3.0.4 for Ubuntu 16.04 and CUDA 9.0 tar package from https://developer.nvidia.com/nvidia-tensorrt-download.
2. Extract archive
```
  tar -xzf TensorRT-3.0.4.Ubuntu-16.04.3.x86_64.cuda-9.0.cudnn7.0.tar.gz
```
3. Install uff python package using pip
```
  sudo pip install TensorRT-3.0.4/uff/uff-0.2.0-py2.py3-none-any.whl
```

Clone and build this project

git clone --recursive https://github.com/NVIDIA-Jetson/tf_to_trt_image_classification.git
cd tf_to_trt_image_classification
mkdir build
cd build
cmake ..
make 
cd ..

Download models and create frozen graphs

Run the following bash script to download all of the pretrained models.

source scripts/download_models.sh

If there are any models you don't want to use, simply remove the URL from the model list in scripts/download_models.sh.
Next, because the TensorFlow models are provided in checkpoint format, we must convert them to frozen graphs for optimization with TensorRT. Run the scripts/models_to_frozen_graphs.py script.

python scripts/models_to_frozen_graphs.py

If you removed any models in the previous step, you must add 'exclude': true to the corresponding item in the NETS dictionary located in scripts/model_meta.py. If you are following the instructions for executing engines below, you may also need some sample images. Run the following script to download a few images from ImageNet.

source scripts/download_images.sh

Convert frozen graph to TensorRT engine

Run the scripts/convert_plan.py script from the root directory of the project, referencing the models table for relevant parameters. For example, to convert the Inception V1 model run the following

python scripts/convert_plan.py data/frozen_graphs/inception_v1.pb data/plans/inception_v1.plan input 224 224 InceptionV1/Logits/SpatialSqueeze 1 0 float

The inputs to the convert_plan.py script are

frozen graph path
output plan path
input node name
input height
input width
output node name
max batch size
max workspace size
data type (float or half)

This script assumes single output single input image models, and may not work out of the box for models other than those in the table above.

Execute TensorRT engine

Call the examples/classify_image program from the root directory of the project, referencing the models table for relevant parameters. For example, to run the Inception V1 model converted as above

./build/examples/classify_image/classify_image data/images/gordon_setter.jpg data/plans/inception_v1.plan data/imagenet_labels_1001.txt input InceptionV1/Logits/SpatialSqueeze inception

For reference, the inputs to the example program are

input image path
plan file path
labels file (one label per line, line number corresponds to index in output)
input node name
output node name
preprocessing function (either vgg or inception)

We provide two image label files in the data folder. Some of the TensorFlow models were trained with an additional "background" class, causing the model to have 1001 outputs instead of 1000. To determine the number of outputs for each model, reference the NETS variable in scripts/model_meta.py.

Benchmark all models

To benchmark all of the models, first convert all of the models that you downloaded above into TensorRT engines. Run the following script to convert all models

python scripts/frozen_graphs_to_plans.py

If you want to change parameters related to TensorRT optimization, just edit the scripts/frozen_graphs_to_plans.py file. Next, to benchmark all of the models run the scripts/test_trt.py script

python scripts/test_trt.py

Once finished, the timing results will be stored at data/test_output_trt.txt. If you want to also benchmark the TensorFlow models, simply run.

python scripts/test_tf.py

The results will be stored at data/test_output_tf.txt. This benchmarking script loads an example image as input, make sure you have downloaded the sample images as above.

NVIDIA-AI-IOT / tf_to_trt_image_classification