TensorRT 教程（二）：TensorRT 源码简介

 

NVIDIA TensorRT是一种高性能神经网络推理(Inference)引擎，用于在生产环境中部署深度学习应用程序，应用有图像分类、分割和目标检测等，可提供最大的推理吞吐量和效率。TensorRT是第一款可编程推理加速器，能加速现有和未来的网络架构。

1. TensorRT 库构成

以编译后源码压缩包安装方式进行安装的TensorRT库主要有以下文件夹：

1 2 3 4 5 6

├── data ├── doc ├── include # 所有头文件，可以查看所有函数的接口和说明 ├── lib # 所有动态链接库.so文件 ├── python # python API 例子 └── samples # c++ 例子

1.1 include文件夹

include文件夹内部包含文件如下：

1 2 3 4 5 6 7 8 9 10 11 12 13

├── include │ ├── NvCaffeParser.h │ ├── NvInfer.h │ ├── NvInferPlugin.h │ ├── NvInferPluginUtils.h │ ├── NvInferRuntimeCommon.h │ ├── NvInferRuntime.h │ ├── NvInferVersion.h │ ├── NvOnnxConfig.h │ ├── NvOnnxParser.h │ ├── NvOnnxParserRuntime.h │ ├── NvUffParser.h │ └── NvUtils.h

下面按照文件之间依赖关系，大致介绍上述头文件。

• NvInferRuntimeCommon.h ：定义了 TRT 运行时所需的基础数据结构（如Dims，PluginField）和大部分基类（class ILogger， class IPluginV2）
• NvInferRuntime.h: 继承NvInferRuntimeCommon.h中的基类，定义了runtime时的拓展功能子类
• NvInfer.h：继承NvInferRuntime.h中的基类，定义了大部分可以直接加入的神经网络层（如class IConvolutionLayer，class IPoolingLayer，class IPoolingLayer）
• NvInferPluginUtils.h：定义plugin layer所需的基础数据结构
• NvInferPlugin.h：初始化注册所有苦衷包含plugin layer
• 其他的从文件名即可看出时分别Caffe，ONNX，UFF 的解析器parser，分别对应训练后的Caffe，Pytorch，TensorFlow网络

1.2 lib以及其他文件夹

NVIDIA官方在 release TRT-7.0 版本时说会放出源码，然而还是以编译后动态链接库的方式放出库文件，所以，老黄的话也不能信，这是人家吃饭的家伙，应该不会放出源码。动态链接库里包含了一些核心的优化代码，比如说如何将 conv+bn+relu 三个层优化成一层，做了哪些工作，这些都在lib文件夹下的.so 里；还有建立整个网络之后使用什么方法寻找一个精度和速度平衡的最终 TRT runtime engine。

samples 下的例子，如果是写 c++ 代码的同学，一定要看看，很多现实中遇到的问题，例子里都有，而且还有一些现成的官方代码，比如说Mask-RCNN帮助写了很多plugin layer，就不用自己实现了。

Python 下的例子包含一些主要用的功能，比如说 parse 网络参数，生成 TRT engine 等。

2. TRT使用流程简介

TRT 大致的工作流程是使用已经训练好的网络参数，将其转换成 TRT engine 序列化代码，供 GPU 进行计算，下图下面将简要介绍三个阶段。

 

2.1 Training

也就是模型训练阶段，这部分如上图所示，目前主流的训练框架都支持。区别只是在从训练框架中提取模型参数权重的方式和所用中间 parser 库不同。

2.2 Developing A Deployment Solution

此阶段就是如何生成一个优化后的 TRT 模型，也是最重要的阶段。这一阶段目标点是，如何在更快地保证 TRT模型运行速度的前提下，保证相对与训练模型精度在可接受范围内。由于模型训练阶段一般使用 float32 精度（当然也有混合精度训练和量化训练），TRT 还支持其他可选项，比如如下图所示的：

• Precision Calibration: TRT 目前支持 INT8 量化
• Layer & Tensor Fusion: 如上文提到的 conv+bn+relu的融合
• Kernel Auto-Tuning: 略
• Dynamic Tensor Memory : 略
• Multi-Stream Execution: 支持多模型的stream流并行处理

 

通过上面的这些步骤，得到优化后的 TRT 模型，就可以通过 serialize 之后保存下来，使得下一次可直接 deserialize 保存下来的 TRT 模型进行 inference。

2.3 Deploying A Solution

将上一阶段序列化之后保存下来的模型，部署到实际生产环境中。由于TRT模型在 Developing A Deployment Solution

是针对平台进行优化的，因此如果在平台A上得到的模型在平台B上不一定适用，所以可以理解为 TRT 针对的是 platform-wise 级别的优化，比如在1080Ti上得到的模型在2080Ti上甚至不能使用。

3. TRT是如何工作的？

TRT 内部有一整套流程来帮助加速 inference，大致内容已经在2.2节中有阐述，有兴趣可看原英文说明。对此，我们并不关心，因为所有这些方法都在未开源的代码中，也没办法自己 hack，所以开始进行 build 之后，我们所有能做的只是等TRT model 的生成。

4. TRT能做些什么？

总的来说，TRT 可以在4个方面帮助开发者进行工作：导入，校准，生成和部署（import, calibration, generate and deploy）模型。

1. 导入模型参数
   TRT 可以直接定义模型，但是我们通常并不这么做，因为对于现有深度学习训练框架，如 Pytorch 和 TensorFlow，我们通常使用对应的 parser 直接获得所需的模型结构和参数，然后进行下一步。TRT 包括的 parser 有如下：Caffe， UFF，ONNX。
2. 配置 optimization profile
   例如是否配置 layer 的动态输入输出。
3. builder
   从我的角度理解，builder 更像是一个外部的配置管理，包括模型的运行资源和 INT8 等配置。
4. engine
   通过 builder 生成的 TRT engine 就是我们最终需要的文件，这个文件可以在 inference 阶段直接载入之后进行运行。

总结

本文介绍了 TRT 库的组成，工作流程和工作原理。

下一篇文章开始，我们将以Pytorch训练得到的 Faster R-CNN 为例，将 TRT 的工作流程展开讲解，方便对各个细节进行研究。