深度学习系列五：使用TensorRT对网络进行加速

概述

TensorRT是NVIDIA一个深度学习加速平台，用于对神经网络模型进行优化，从而加速其推理过程，现已被广泛应用于嵌入式芯片和智能汽车平台等，提高其GPU的推理速度，以实现更快的响应速度或者降低平台的性能消耗。TensorRT目前已经对Tensorflow、Pytorch、ONNX、Caffe等多个深度学习框架有了良好的支持，在最新的Tensorflow2.0中，TensorRT5.0已经被集成在其中，对于Tensorflow用户而言，可以方便的使用TF-TFT对已有的模型进行相应的优化，但由于一些原因，本人并未使用TF-TRT。

对于其工作原理有兴趣的，可以直接参考Speed up TensorFlow Inference on GPUs with TensorRT，该文章有较为详细的分析，在此不多作赘述。

TensorRT的推断代码是直接利用cuda语言运行在显卡上的，所有的代码库仅仅包括C++和cuda，上层也有python的封装，提供一些供python使用的api，当我们利用这个库运行我们训练好的模型时，其运行速度和所占内存的大小都会大大缩减。使用TensorRT时，需要对cuda的使用有一点最基础的了解。

此外，深度学习项目在训练时为了加快速度或者各种其他原因（使用多GPU分布式训练等），会安装比较复杂的环境。但在实际部署时，却可能仅仅使用单GPU嵌入式平台（例如 NVIDIA Jetson）或者在云端进行部署，此时，部署端往往也要有与训练时相同的深度学习环境，这对于部署而言，是一件繁杂的任务。利用TensorRT进行部署，不仅可以加快网络推理速度，而且可以将训练好的网络进过适当的格式转换后直接丢进TensorRT进行运行，而不再依赖原来的训练环境，可谓一举多得。

安装

以TensorRT6.0.1为例，来源于官网安装教程

主要分为如下几步：

• 下载
  官方下载地址：https://developer.nvidia.com/tensorrt.
  选择合适的版本下载，这里以6.0 GA，Debian为例，环境为ubuntu16.04，下载时一定需要确定自己的cuda和驱动安装完好，且选择对应的版本进行下载。
• 安装
  运行如下几条指令：

$ sudo dpkg -i  
nv-tensorrt-repo-ubuntu1x04-cudax.x-trt6.0.x.x-ga-yyyymmdd_1-1_amd64.deb
$ sudo apt-key add /var/nv-tensorrt-repo-cudax.x-trt6.0.x.x-ga-yyyymmdd/7fa2af80.pub

$ sudo apt-get update
$ sudo apt-get install tensorrt

对于Python 2.7:

$ sudo apt-get install python-libnvinfer-dev

对于Python 3.x:

$ sudo apt-get install python3-libnvinfer-dev

如果基于TensorFlow作为训练环境:

$ sudo apt-get install uff-converter-tf

• 检验

$ dpkg -l | grep TensorRT

安装正常将会看到：

ii  graphsurgeon-tf                                             6.0.1-1+cuda10.0                             amd64        GraphSurgeon for TensorRT package
ii  libnvinfer-bin                                              6.0.1-1+cuda10.0                             amd64        TensorRT binaries
ii  libnvinfer-dev                                              6.0.1-1+cuda10.0                             amd64        TensorRT development libraries and headers
ii  libnvinfer-doc                                              6.0.1-1+cuda10.0                             all          TensorRT documentation
ii  libnvinfer-plugin-dev                                       6.0.1-1+cuda10.0                             amd64        TensorRT plugin libraries
ii  libnvinfer-plugin6                                          6.0.1-1+cuda10.0                             amd64        TensorRT plugin libraries
ii  libnvinfer-samples                                          6.0.1-1+cuda10.0                             all          TensorRT samples
ii  libnvinfer6                                                 6.0.1-1+cuda10.0                             amd64        TensorRT runtime libraries
ii  libnvonnxparsers-dev                                        6.0.1-1+cuda10.0                             amd64        TensorRT ONNX libraries
ii  libnvonnxparsers6                                           6.0.1-1+cuda10.0                             amd64        TensorRT ONNX libraries
ii  libnvparsers-dev                                            6.0.1-1+cuda10.0                             amd64        TensorRT parsers libraries
ii  libnvparsers6                                               6.0.1-1+cuda10.0                             amd64        TensorRT parsers libraries
ii  python3-libnvinfer                                          6.0.1-1+cuda10.0                             amd64        Python 3 bindings for TensorRT
ii  python3-libnvinfer-dev                                      6.0.1-1+cuda10.0                             amd64        Python 3 development package for TensorRT
ii  tensorrt                                                    6.0.1.5-1+cuda10.0                           amd64        Meta package of TensorRT
ii  uff-converter-tf                                            6.0.1-1+cuda10.0                             amd64        UFF converter for TensorRT package

此外，如果需要在python中使用，还需用安装pycuda：

pip install 'pycuda>=2017.1.1'

至此所有的安装工作已经完成，有任何问题可以参考以上官网教程。

使用

本人训练平台为Tensorflow2.0，但由于目前最新的TensorRT6.0.1在转换模型中应用到Tensorflow相关库，且目前只支持到1.14，因此在转换过程中将Tensorflow降级为1.14，转换完成后，运行实际推理将不再依赖Tensorflow。此外前文提到过，对于Tensorflow用户，本可以直接使用tf-trt进行转换和推理，但由于本人Tensorflow版本为2.0，该版本只支持到tensorrt5.0版本，对于模型中的部分层无法支持转换，且tensorflow2.0已经在大力推行其集成的keras上层接口，相应的1.0底层接口在未来可能将被遗弃，而使用tf-trt的相关文档和资料都需要用tensorflow底层接口实现，因此思虑再三，还是放弃了使用tf-trt，而直接使用tensorrt进行推理，待tf-trt进一步完善后再继续跟进。

• 1.转换模型
  首先需要将tensorflow模型转换为.pb模型，目前TensorRT仅支持对于该格式的转换，熟悉tensorflow的应该都知道如何输出pb模型，这里不多赘述，不会的可以参考Tensorflow 模型保存与恢复（3）保存模型到单个文件中。
  有了pb模型之后，可以利用TensorRT提供的转换工具将该模型转换为uff格式模型:

python3.5 /usr/lib/python3.5/dist-packages/uff/bin/convert_to_uff.py --input_file models/xx.pb

将会在同文件夹下生成uff文件，此外，当有部分层无法支持时，也需要手动加载自定义plugin并利用相应API进行转换，否则生成的uff文件是无法运行的，该部分将在后续文章中详细描述。

• 2.建立引擎
  有了uff文件之后，需要建立一个TensorRT引擎，该引擎包含了整个优化后的推理模型、相应绑定的输入输出以及一些配置，一旦有了这个引擎之后，就可以在TensorRT独立运行而不再依赖其训练环境了。
  代码参考：

import tensorrt as trt
with builder = trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.UffParser() as parser:
    	parser.register_input("Placeholder", (1, 28, 28))
    	parser.register_output("fc2/Relu")
		parser.parse(model_file, network)

其中register_input即注册网络的输入，register_output即注册网络的输出，输入输出名字必须与网络原有的一致，同时必须指定输入维度信息，如果存在多个输入输出，需要全部依次注册。
接下来就可以创建引擎了：

builder.max_batch_size = max_batch_size
builder.max_workspace_size = 1 <<  20 #指定分配给引擎的显存
with trt.Builder(TRT_LOGGER) as builder:
	with builder.build_cuda_engine(network) as engine:

创建引擎之后，强烈建议将引擎保存下来:

serialized_engine = engine.serialize()
with open(engine_path, 'wb') as f:
                f.write(serialized_engine)

有了引擎文件之后，就不需要再重复创建引擎了，只需要加载即可：

with open(engine_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
                return runtime.deserialize_cuda_engine(f.read())

• 3.开始推理
  用TensorRT进行推理的过程相对麻烦一点，需要将输入数据按照注册的格式序列化到连续内存中，再将该片数据拷到预先分配的内存中，最后由GPU进行计算，然后将结果返回，因此，需要一些额外内存分配操作来完成以上流程：

import pycuda.driver as cuda
		h_input = cuda.pagelocked_empty(engine.get_binding_shape(0).volume(), dtype=np.float32)
		h_output = cuda.pagelocked_empty(engine.get_binding_shape(1).volume(), dtype=np.float32)
		# Allocate device memory for inputs and outputs.
		d_input = cuda.mem_alloc(h_input.nbytes)
		d_output = cuda.mem_alloc(h_output.nbytes)
		# Create a stream in which to copy inputs/outputs and run inference.
		stream = cuda.Stream()
		with engine.create_execution_context() as context:
			# Transfer input data to the GPU.
			cuda.memcpy_htod_async(d_input, h_input, stream)
			# Run inference.
			context.execute_async(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
			# Transfer predictions back from the GPU.
			cuda.memcpy_dtoh_async(h_output, d_output, stream)
			# Synchronize the stream
			stream.synchronize()
			# Return the host output. 
return h_output

以上为官方的示例代码，基本可以满足使用了，需要注意的是，如果存在多输入多输出时，需要给每个输入和输出都预先分配内存和显存，在提供输入时，也需要把多输入都填充，输出也会按照注册时的输出数量输出多个输出。

总结

至此，整个TensorRT的基本使用方法就差不多了，TensorRT安装后位于/usr/src/tensorrt/sample中有大量的官方示例供参考，可以帮助解决很多实际使用中的问题，此外还有一些例如使用plugin和创建INT8类型以及校准引擎的方法会在后续文章中记录。