基于 TensorRT 的模型推理加速

1 TensorRT 概述

       TensorRT 是由英伟达公司推出的一款用于高性能深度学习模型推理的软件开发工具包，该软件开发工具包中包含了应用程序运行时环境和深度学习模型推理优化器，可以把经过优化后的深度学习模型构建成推理引擎部署在实际的生产环境中，其主要的应用场景包括计算机视觉各个领域和推荐系统等。TensorRT 是采用 C++语言编写的高性能推理框架，可以在所有平台上提供 C++语言的实现，并且也能在 Linux 平台上提供 Python 语言的实现。同时，TensorRT 推理应用程序与英伟达的 GPU 结合可以提供最大的吞吐量和较低的延迟速度。 

      现在大多数的深度学习网络模型结构复杂并且参数量巨大，需要使用多个高性能的 GPU 分布式训练才能获得全局最优的结果，这使得深度学习方法很难应用中在实际的生产环境中。主要原因如下：为了降低生产成本，应用程序通常都会部署在嵌入式开发板上，其算力相对较低，对于结构复杂且参数量巨大的网络模型较难获取实时的推理速度；在使用深度学习方法时，部署应用的机器的深度学习环境要与网络模型训练时的环境相同，这增加了深度学习方法部署的复杂性。与许多轻量级的深度学习网络模型相比，TensorRT 不仅可以大幅提高网络的推理速度，而且只损失些许的精度。同时，使用 TensorRT 优化后的网络模型的部署不再需要与训练时相同的环境。图 1 是 TensorRT 对 ResNet50 优化之后的推理速度对比，从图中可以看到，使用 TensorRT 优化之后的网络推理速度有了较大的提高。

 

      目前，TensorRT 已经支持对常用的深度学习框架（例如，PyThorch、TensorFlow 等）训练出来的网络模型进行优化。同时，TensorRT 推理框架还不够完善，仅支持常见的一些操作，比如 ReLU、SoftMax、全连接、池化和卷积等；对于不支持的操作，TensorRT 支持用户自定义网络模型结构。使用 TensorRT 优化并生成的推理引擎可以应用在不同的 GPU 和嵌入式开发平台上，例如 TESLA P4、TESLA P100、DRIVE PX 2、JETSON TX2  和 NVIDIA DLA 等。 

 

      TensorRT 的使用流程如上图所示，主要包括构建推理引擎和应用推理引擎两个阶段。构建推理引擎阶段是使用 TensorRT 推理优化器对已经训练好的深度学习网络模型结构及模型参数进行优化。通常情况下，推理引擎的构建过程将花费很长的时间，尤其是在一些嵌入式开发平台上。因此，我们通常是先单独构建推理引擎，然后将其序列化为二进制文件并存储在磁盘上，在应用程序中使用时只需要反序列推理引擎文件即可快速地构建推理引擎，节约了大量的时间。同时，这样的使用方式也带来了另一个好处，对网络模型结构的修改只需要重新生成推理引擎二进制文件即可，无须修改应用程序源代码。应用推理引擎阶段是反序列化推理引擎二进制文件生成运行时推理引擎，然后就可以输入图像数据并进行模型推理，最终获得输出结果。 

2 TensorRT 性能优化方式

     对于训练好的深度学习网络模型结构和权重参数，在导入到 TensorRT 推理引擎优化器之后，会对其在模型结构和权重参数上进行一系列的全局优化，优化过程如图 3 所示。

 

     TensorRT 对网络模型结构和权重参数的优化主要包括层与张量融合和权重与激活精度校准，本节将对这两种优化方式进行简单地介绍。 

2.1  层与张量融合

      对于实例分割领域来说，为了追求更高的检测框平均精度和分割掩码平均精度，其网络模型结构越来越复杂，包含大量的卷积层和全连接层等神经网络模块。在对网络模型做前向推理时，网络结构中的每一个卷积操作都需要 GPU 调用不同的 CUDA 核心去执行计算来获取结果。对于卷积操作中的张量运算，CUDA核心的计算速度是非常快的，但是，大量的卷积层和全连接层会使得 CUDA 核心的调用和张量的读写浪费大量的时间。同时，主存和显存之间大量的数据拷贝操作也是对 GPU 时间的浪费。TensorRT 通过对不同网络层之间的垂直融合和水平融合，使得网络结构的层数大大减少。同时，也会删除没有使用到的输出层以及一些等同没有操作的层结构。 

       如图 4 所示，层间垂直融合就是把卷积层、偏置层和激活层融合成一个CBR 网络模块结构，通过这种方式的层间垂直融合之后，在网络模型执行前向推理时，GPU 只需要调用一个 CUDA 核心。 

      如图5 所示，层间水平融合就是把多个权重参数不同且结构相同层融合成一个更宽的层结构，通过这种方式的层间水平融合之后，在网络模型执行前向推理时，GPU 只需要调用一个 CUDA 核心。 

     通过对庞大且复杂的深度神经网络结构进行层间水平融合和层间垂直融合之后，网络结构中的层数更少，在前向推理的过程中，使用到的 CUDA 核心数量也就更少，进而降低了 CUDA 核心调用和张量读写的时间并提高了 CUDA 核心的利用率，因此整个深度神经网络的结构更小，推理速度更快。 

2.2  权重与激活精度校准

      在深度神经网络模型训练的后期阶段，由于反向传播梯度变化很小，如果权重参数的精度不够，训练结果将无法进一步收敛，因此，许多深度学习框架中的张量存储和运算过程都是使用 32 位的浮点数（Full 32-bit precision，FP32）。然而，在深度神经网络模型训练完成之后，在应用部署的程序中不需要反向传播，只有前向推理的过程，因此，可以通过降低数据精度的方式来提高网络前向推理的速度以及降低显存的空间占用。但是，在使用低精度表示的数据进行前向推理时会存在信息丢失的问题，从而导致推理结果性能的下降。目前，TensorRT 已经支持 FP16 和 INT8 的计算，在对深度神经网络模型结构的优化的过程中，TensorRT 也会对权重参数进行优化校准，使得推理结果性能的损失降到最低。从图 6 中可以看出，TensorRT 使用 INT8 精度的数据获取的前向推理的结果仅损失了不到 0.5%的性能。