Introduction

TensorRT is a Software Development Kit(SDK) from Nvidia. It enables developers to develop highly efficient neural network models, which are capable of providing signficantly low inference time than their non-trt optimized counterparts. As useful as this technology is, it is not the easiest one to install and get running, given a plethora of it's dependencies. In this TensorRT tutorial we are going start off by installing TensorRT in a conda environment, as it is the most easy and reliable way to get it up and running quickly, after which we will look into optimizing a model using TensorRT. Let's begin!

Requirements

It is assumed that following requirements are met, before proceeding with any of the below installation.

• Ubuntu 20.04 LTS.
• Nvidia Graphics Card, with CUDA compatible display drivers.
• GCC compiler version 7.x or above.
• Anaconda

Installation and Configuration

Let us start with creating the conda environment. It can be done easily by typing in the following command into terminal. Replace "your-enviroment-name" by the name you prefer for this env.

                            
                                $ conda create -n "your-enviroment-name" tensorflow-gpu==2.4.1
                            

The above command creates a new conda environment and installs tensorflow-gpu on the latest python version available. However you can specify the python version by adding "python==3.x" in the end like so, and activating it.

                                
                                    $ conda create -n "your-enviroment-name" tensorflow-gpu==2.4.1 python==3.8
                                    $ conda activate "your-enviroment-name"
                            

We can verify the installation of tensorflow-gpu at this point, by importing tensorflow in python and checking if gpu is visible to tensorflow, as follows.

                            
                                $ python
                                >>> import tensorflow as tf
                                >>> tf.__version__
                                >>> tf.test.is_gpu_available() # this should return true
                                >>> quit()
                            

[Optional] Install other libraries like matplotlib, pandas, scikit-learn, etc.. as follows.

                            
                                $ conda install -c anaconda pandas matplotlib scikit-learn
                            

Installing TensorRT in conda environment.

                            
                                $ pip install nvidia-pyindex
                                $ pip install nvidia-tensorrt
                            

After this installation the compiler will change from "[GCC 7.3.0] :: Anaconda, Inc. on linux" (version may differ) to something like "[GCC 7.5.0] on linux" depending on the compiler installed on your system. You can check the compiler, by simply typing python in the terminal and pressing enter. An output like this should show up.

Optimizing a deep learning model using TensorRT

We can now verify the TensorRT installation by optimizing a model. I have used Xception Net here, but you can use any model as long as it is in .pb (protobuf) format, by just changing the path to your model in the following python script.

                                
                                from tensorflow.python.compiler.tensorrt import trt_convert as trt

                                input_saved_model_dir = "./Reference_Data/Model/Xception/"
                                output_saved_model_dir = "./Reference_Data/Model/Xception/TRT_Model"

                                print("Conversion : Started")
                                converter = trt.TrtGraphConverterV2(input_saved_model_dir=input_saved_model_dir)
                                converter.convert()
                                converter.save(output_saved_model_dir)
                                print("Conversion : Finished")
                            

The Output of the above code should look like this in jupyter notebook or any other IDE.

                                
                                    Conversion : Started
                                    INFO:tensorflow:Linked TensorRT version: (0, 0, 0)
                                    INFO:tensorflow:Loaded TensorRT version: (0, 0, 0)
                                    INFO:tensorflow:Assets written to: ./Reference_Data/Model/Xception/TRT_Model/assets
                                    Conversion : Finished