TensorRT Inference引擎简介及加速原理简介

简介

最近在做CNN卷积神经网络量化方面的工作，查阅资料发现TensorRT有新颖的思想，记录学习的知识，如有问题请指教！TensorRT是NVIDIA 推出的一款基于CUDA和cudnn的神经网络推断加速引擎(C++库)。相比于一般的深度学习框架，在CPU或者GPU模式下其可提供10X乃至100X的加速，极大提高了深度学习模型在边缘设备上的推断速度。TensorRT项目立项的时候名字叫做GPU Inference Engine（简称GIE）。目前TensorRT已经支持Caffe、Caffe2、TensorFlow、MxNet、Pytorch等主流深度学习库。

TensorRT加速原理

在计算资源并不丰富的嵌入式设备上，TensorRT之所以能加速神经网络的的推断主要得益于两点。首先是TensorRT支持INT8和FP16的计算，通过在减少计算量和保持精度之间达到一个理想的trade-off，达到加速推断的目的。

更为重要的是TensorRT对于网络结构进行了重构和优化，主要体现在一下几个方面：

第一是tensorRT通过解析网络模型将网络中无用的输出层消除以减小计算。

第二是对于网络结构的垂直整合，即将目前主流神经网络的conv、BN、Relu三个层融合为了一个层，例如将图1所示的常见的Inception结构重构为图2所示的网络结构。

图1. Inception结构

图2. Inception结构重构 第三是对于网络的水平组合，水平组合是指将输入为相同张量和执行相同操作的层融合一起，如图2向图3的转化。

图3. 相同操作的层水平融合

TensorRT直接支持的层

层名称	说明
Convolution	Convolution层计算3D卷积（通道，高，宽），可选择是否使用偏置。
Pooling	Pooling层在通道维度实现池化操作。支持最大池化与均值池化。
Activation	激活层通过逐像素激活函数实现。支持的激活类型包括，RELU，tanh与sigmoid。
Concatenation	连接层在通道维度上将宽高相同的tensors进行连接。
FullConnected	全连接层实现了矩阵向量积运算，可选择是否使用偏置。
Scale	Scale层实现了逐tensor，逐通道，逐权重的仿射变换，与或运算，以及常数幂运算。
LRN	LRN层实现跨通道局部响应归一化。

注：BatchNormalization可以利用TensorRT的Scale层实现。

TensorRT–8-bit Inference

流程图：

实现的基本思路是首先构建一种FP32数据向INT8数据的映射关系，该映射中的边界并不是两种数据类型的最大值，而是将FP32设置成一个Threshold，将这个阈值与INT8的最大值(127)构建映射关系，具体实现形式是通过一个scale来进行对应。

首先的问题是，这种映射关系的阈值如何寻找呢？不同的网路显然这一阈值是不同的，因此我们需要一个矫正数据集(calibration dataset)来进行scale的选取，其选择的标准为最小化KL_divergence，P, Q -两个离散概率分布：

简单来说，因为FP32～int8只是重新编码信息，所以选择阈值时应尽量保证减少信息的丢失。
其次的问题是scale的表示与获取。scale的表示相对而言较为简单，如下形式：

然后在TensorRT工作流运行时运行矫正算法，来优化scale系数。

如何矫正呢？首先，我们在矫正数据集上运行FP32精度的inference, 对于网络的每一层来说，收集激活时的直方图，从而产生不同saturation thresholds对应的量化分布图，然后选取能够最小化KL_divergence(ref_distr, quant_distr)的阈值作为选择的阈值。

确定阈值后，进行8bit量化，如下图，需要考虑溢出问题。

结果