Installing TensorRT in Jetson TX2
TensorRT is an optimization tool provided by NVIDIA that applies graph optimization and layer fusion, and finds the fastest implementation of a deep learning model. In other words, TensorRT will optimize our deep learning model so that we expect a faster inference time than the original model (before optimization), such as 5x faster or 2x faster. The bigger model we have, the bigger space for TensorRT to optimize the model. Furthermore, this TensorRT supports all NVIDIA GPU devices, such as 1080Ti, Titan XP for Desktop, and Jetson TX1, TX2 for embedded device.
There are at least two options to optimize a deep learning model using TensorRT, by using: (i) TF-TRT (Tensorflow to TensorRT), and (ii) TensorRT C++ API. In this post, we will specifically discuss how we can install and setup for the first option, which is TF-TRT, in Jetson TX2. Here is the step-by-step process:
- Flashing the Jetson TX2
By flashing the Jetson TX2 using Jetpack, it will already install CUDA, CuDNN, TensorRT, etc . You can follow the video below to flash your TX2.
- Installing Python 3 and PIP 3
$ sudo apt-get install -y python3-pip python3-dev
- Installing Tensorflow
Download the Tensorflow wheel file here, and then install it by using:sudo pip3 install your_wheel_file.whl
For how we can optimize a deep learning model using TensorRT, you can follow this video series here: