Skip to content

Python sample for referencing pre-trained SSD MobileNet V2 (TF 1.x) model with TensorRT

Notifications You must be signed in to change notification settings

brokenerk/TRT-SSD-MobileNetV2

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

11 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

TensorRT Python Sample for a Re-Trained SSD MobileNet V2 Model (only faces' detection)

Original NVIDIA sample's GiHub repository: AastaNV/TRT_object_detection


Tested on a NVIDIA Jetson AGX Xavier with Jetpack 4.3 and Tensorflow 1.15.


Performance includes memcpy and inference.

Model Input Size TRT Nano
ssd_mobilenet_v2_coco 300x300 46ms

Since the optimization of preprocessing is not ready yet, image read/write time is not included here.

Install Tensorflow 1 dependencies and PyCUDA

$ sudo apt-get update
$ sudo apt-get install libhdf5-serial-dev hdf5-tools libhdf5-dev zlib1g-dev zip libjpeg8-dev liblapack-dev libblas-dev gfortran
$ sudo apt-get install python3-pip
$ pip3 install -U pip testresources setuptools==49.6.0
$ pip3 install -U numpy==1.16.1 future==0.18.2 mock==3.0.5 h5py==2.10.0 keras_preprocessing==1.1.1 keras_applications==1.0.8 gast==0.2.2 futures protobuf pybind11
$ pip3 install --pre --extra-index-url https://developer.download.nvidia.com/compute/redist/jp/v43 'tensorflow<2'
$ pip3 install numpy pycuda --user


Prepare your pre-trained model

The base object detection model is available here: TensorFlow model zoo.
Remember that this sample is adjusted only for re-trained SSD MobileNet V2 models (use the frozen_inference_graph.pb file, exported after your custom training).
If original sample is required, visit: AastaNV/TRT_object_detection

$ git clone https://github.com/brokenerk/TRT-SSD-MobileNetV2.git
$ cd TRT-SSD-MobileNetV2
$ mkdir model
$ cp [model].tar.gz model/
$ tar zxvf model/[model].tar.gz -C model/
// ============================================================================
// Or just put your frozen_inference_graph.pb file inside the model/ directory
// ============================================================================
Supported models:
  • ssd_mobilenet_v2_coco


Update graphsurgeon converter

Edit /usr/lib/python3.6/dist-packages/graphsurgeon/node_manipulation.py

def create_node(name, op=None, _do_suffix=False, **kwargs):
     node = NodeDef()
     node.name = name
     node.op = op if op else name
     node.attr["dtype"].type = 1
     for key, val in kwargs.items():
         if key == "dtype":
             node.attr["dtype"].type = val.as_datatype_enum


RUN

1. Maximize the Nano performance

$ sudo nvpmodel -m 0
$ sudo jetson_clocks

2. Execute

$ python3 main.py [test_image_path]

It takes some time to compile a TensorRT model when the first launching.
After that, TensorRT engine can be created directly with the serialized .bin file

@ To get more memory, it's recommended to turn-off X-server.




About

Python sample for referencing pre-trained SSD MobileNet V2 (TF 1.x) model with TensorRT

Topics

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages