OCR技术专题系列二：TensorRT入门实践

一、简介

我们在Tensorflow、Caffe、Pytorch等训练框架上训练神经网络之后，需要对模型迁移到部署框架上进行部署。TensorRT是Nvidia开发的一个神经网络前向推理加速的C++库，用户无需像剪枝那样在训练时对模型进行定制化处理，只需把模型提供给TensorRT即可实现加速。TensorRT支持多种模型的解析，包括:Tensorflow pb->uff->TensorRT, Pytorch pth-> onnx-> TensorRT等转换流程。

二、基本原理

TensorRT对模型的优化包括：

神经网络模型计算图优化。包括合并convolution, Bias和ReLU为计算单元CBR、GPU并发、消除concat等。

将FP32转为INT8量化、FP16低精度运算等。

自动选取最优CUDA kernel。矩阵乘法、卷积有多种CUDA实现方式，TensorRT根据数据大小和形状自动选取最优实现。

三、安装和使用

目前TensorRT提供了C++与Python的API接口，本文使用Python接口介绍TensorRT框架的一般使用流程。

(a)下载：登录官网地址[1]，根据操作系统、cuda版本选择适合自己的版本。我这里使用tar包安装，选择Ubuntu16 cuda10，下载的文件是TensorRT-7.0.0.11.Ubuntu-16.04.x86_64-gnu.cuda-10.0.cudnn7.6.tar.gz

图1[1]

(b)安装:参考官方文档[2]。主要包括TensorRT, uff, graphsurgeon三个whl的安装以及pycuda的安装。下面是详细流程。

解压:

tar xzvf TensorRT-7.0.0.11.Ubuntu-16.04.x86_64-gnu.cuda-10.0.cudnn7.6.tar.gz

安装TensorRT：

pip install tensorrt-7.0.0.11-cp35-none-linux_x86_64.whl

安装graphsurgeon

pip install graphsurgeon-0.4.1-py2.py3-none-any.whl

安装UFF：

pip install uff-0.6.5-py2.py3-none-any.whl

检验是否可以将pb模型转换为uff格式（需要先安装好tensorflow-gpu）：进入目录data/mnist下，有lenet5_mnist_frozen.pb模型文件：

cd TensorRT-7.0.0.11/data/mnist/

使用命令转换得到lenet5_mnist_frozen.uff：

convert-to-uff lenet5_mnist_frozen.pb

如果出现“ImportError: libcublas.so.X.0: cannot open shared object file: No such file or directory”之类的报错，可能是公共服务器上有多个cudnn和cuda版本，需要自己使用的时候export匹配版本的LD_LIBRARY_PATH，参见(c)。

安装pycuda

pip install pycuda

(c)export相关路径

将下列${path}的路径换成自己的就好。

tensorRT的路径：

export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${ tensorRT path}/lib

cuda和cudnn相关：

export LD_LIBRARY_PATH= ${cuda path}/lib64:${LD_LIBRARY_PATH}

export LD_LIBRARY_PATH=${ cudnn path}/lib64/:${ cudnn path}/include:${LD_LIBRARY_PATH}

export PATH= ${cuda path}/bin/:${PATH}

(d)测试

解压缩后的tensorRT目录下有samples文件夹，包含一些测试用例（c++和python的）。其中目录samples/python/是python的测试用例。

这里介绍samples/python/ end_to_end_tensorflow_mnist的使用。

几个注意点：

推理代码使用的模型是lenet.uff，在data/mnist/下。进入data/mnist/下运行“python download_pgms.py”可下载数据集，得到的文件格式为*.pgm。路径设置示例如图2所示。

图2

官方代码的模型输入和输出节点名称是错误的，使用netron打开模型查看即可发现。我这里把它们改成了in, out。

图3

设置完这些，python sample.py即可。

四、主要代码分析

(a)设置输入尺寸、输入节点和输出节点名称，如图3所示。

(b)创建cuda engine：新建 parser注册输入输出节点名称和输入尺寸，并解析uff文件到network中，最后用builder创建network对应的engine。默认使用FP32加速，即只是图优化等操作，不进行量化或降精度。设置FP16十分简单，一行代码即可：

builder.fp16_mode = True

FP32转INT8对模型性能影响较大，需要校准集校准过程，这里不做介绍。

图4

Cuda engine的建立比较花费时间，第一次新建后可以序列化保存到本地，下次直接load即可（图5）。

图5

(c) 构建输入输出缓冲区：binding包括输入和输出节点

图6

(d)预处理和拷贝输入数据：原模型读取数据以及预处理的代码，应在np.copyto()方法前执行。注意：此处raval()将输入数据摊平，默认摊平前图像为NCHW格式，弄错四个维度的顺序（如NHWC）会导致输出结果错误。

图7

(e)推理和搬数据

图8

整体的Main函数流程如下：

图9

[1]https://developer.nvidia.com/nvidia-tensorrt-7x-download

[2]https://docs.nvidia.com/deeplearning/tensorrt/archives/tensorrt-700/tensorrt-install-guide/index.html#installing-tar