TensorRT部署YOLOv5(03)-TensorRT介绍

TensorRT是本专栏中最重要的内容，绝大多数内容将围绕TensorRT来展开，本文对TensorRT进行一个基本的介绍，让不熟悉TensorRT的读者能够对TensorRT是什么，如何使用它有一个较为全面的认识

 

Nvidia TensorRT是一个用于Nvidia GPU上高性能机器学习推理的SDK，对开发者屏蔽了模型在GPU上推理运行的CUDA计算细节，用户只需要通过一套简介易用的Python/C++ API接口，即可方便的将模型在GPU上进行加速推理。另外TensorRT还提供了模型转换、性能评估的工具，方便用户将各训练框架生成的模型转换到TensorRT能够识别的形式，以及在未进行业务开发之前快速评估模型推理的性能

Nvidia官方网站的TensorRT介绍页面有一张TensorRT的工作流程示意图

 这张图非常清晰的解释了TensorRT的工作，预训练网络(Trained Neural Network)在经过精度校准(Precision Calibration)、层融合(Layer Fusion)、算子自动调谐(Kernel Auto-Tuning)、动态张量内存(Dynamic Tensor Memory)、多流执行(Multi-Stream Execution)，最终生成预优化的推理引擎(Optimized Inference Engine)

Build Phase and Runtime Phase

TensorRT的主要工作内容可分为构建期(Build Phase)和运行期(Runtime Phase)

构建期的目标是生成一个能够被运行期加载并运行的TensorRT推理引擎(Engine)，模型的来源可以是从各种深度学习框架(TensorFlow、Pytorch、Caffe等)训练好的模型文件，也可以是利用TensorRT API原生搭建的网络，从深度学习框架导出的权重文件中加载权重参数。推理引擎是构建期的输出，推理引擎可以以本地二进制文件的形式存在，也可以是程序中的一个类实例。构建期在引擎的创建过程中，除了进行模型结构的解析以及生成引擎之外，还有很多中间的优化工作，例如计算图优化，经典的操作Conv+Add+ReLU层融合，节点消除，节点变换(Pad、Slice、Concat、Shuffle)，并对算子在GPU上的实现进行本地运行评估和选择，所有这些优化都是工具自动进行的，但是用户可以通过一些额外的参数来调整优化过程。经过构建期生成的推理引擎，其内部的网络结构和原模型已经完全不同

运行期的作用是加载推理引擎，并在GPU上进行执行，在这个阶段，用户需要为引擎提供输入数据，并准备好输出内存，通过TensorRT提供的运行期API进行模型执行

构建期和运行期可以处于同一个程序中，也可以分开独立进行。例如可以在同一个程序中，先进行构建期将模型转换为推理引擎，然后进行执行期将推理引擎的类实例直接运行推理计算；也可以在一个单独的构建程序中进行模型转换，生成推理引擎，通过API导出为序列化的引擎文件，然后在另一个推理程序中加载引擎文件，生成引擎类实例，进行推理计算。由于模型到引擎的转换会经过比较漫长的优化调谐过程，因此不建议将构建期和运行期放在同一个程序中运行，而比较推荐一次单独构建并导出引擎文件，后续的多次运行都加载同一个引擎文件，这样节省了很多优化耗时。本文后续将以这种方式介绍为主

Workflow

在介绍使用TensorRT进行模型推理的基本工作流程之前，首先需要介绍一下构建TensorRT引擎的几种方式，不同方式的workflow不太相同，而本文主要介绍我自己使用的workflow。在构建期，TensorRT提供了3种从模型构建引擎的方式

• 框架自带的TRT接口
• 使用Parser
• TensorRT API原生搭建

这三种方式各有优劣，下面对这三种方式进行介绍

框架自带TRT接口

框架自带TRT接口指的是深度学习框架例如用TensorFlow、Pytorch等，Nvidia和部分框架有合作，这些框架内集成了TensorRT的部分功能，例如TensorFlow的TF-TRT、Pytorch的Torch-TensorRT，使用这些框架内置的TensorRT接口，可以将框架训练好的模型无缝衔接到TensorRT，直接在现有框架上调用TensorRT相关接口就可以方便的进行模型转换和推理。这种方式的优势就是非常方便，环境统一，开发效率较高，需要做的改动也最少，算子支持的也最好，遇到TensorRT不支持的算子会返回原框架进行计算。缺点也非常明显，首先框架适配的TensorRT是阉割版的TensorRT，不能最大限度的利用TensorRT的优化加速能力，其次TensorRT无法和所有的框架厂商都进行合作与适配，这导致使用其他框架的用户无法使用这种方式。另外，在很多资源有限的嵌入式设备上安装Tensorflow这种较大的深度学习工具也不现实，因此这种方式通常是应用于服务端推理。本文不对这种方式进行介绍

使用Parser

使用Parser方式，指的是将深度学习框架导出的模型文件，经过一个中间表示，转换到TensorRT引擎，中间表示主要是使用ONNX，这种方式的优点是TensorRT能够在构建期尽可能多的操作网络结构和算子优化，因此推理性能较高。缺点也是非常明显的，由于使用了ONNX，原框架中支持的算子与ONNX支持的算子以及TensorRT支持的算子，这三者之间并不是完全覆盖的，在网络模型使用了较多较新层结构时，在算子支持方面可能会存在较多问题，并且由于ONNX本身对不支持的算子也会进行模型结构的变换，这部分并不可控，对性能也有一些损失，这种情况下要么修改网络结构适配ONNX，要么以插件的形式通过TensorRT提供的API手写自定义算子，比较复杂

TensorRT API原生搭建

TensorRT API本身提供了构建网络结构的API，这种方式的优点是网络细节完全由用户控制，某些情况下TensorRT转换出来的网络结构可能并非最佳方案，手工设置的网络结构性能更佳，这种情况下由用户自己搭建TensorRT的原生网络性能能够达到最大化，但是由于TensorRT API搭建网络也存在算子支持的问题，并且从网络整体结构层面进行性能优化本身是一个难度很高的事情，需要用户对计算图优化有较为深入的理解，因此这种方式开发难度最大

总的来说，无论是哪种方式，特殊算子适配是一个绕不开的问题，必须研究插件写法以及CUDA计算细节。本文由于重点是介绍TensorRT上运行YOLOv5模型的全流程，重点将会放在整个流程的完整性上，且YOLOv5模型本身使用到的层和算子比较常规，不存在算子适配问题，因此后文主要以“使用Parser”方式进行介绍，对插件写法和TensorRT API搭建网络等方面暂不进行介绍

以下给出本文实验使用Parser进行模型推理的基本流程

• 导出预训练模型

• 预训练模型转换为ONNX模型

• ONNX模型调整

• ONNX模型转换为TensorRT引擎

  • trtexec

  • Python/C++ API

• 加载TensorRT引擎并进行推理

下图是本实验使用的workflow示意图

​​​​​​​

 在Windows系统上使用TensorFlow进行了YOLOv5的模型训练，导出model.pb模型文件；将model.pb模型文件拷贝到Ubuntu虚拟机中，使用虚拟机中安装的tf2onnx工具将model.pb转换为model.onnx(还需要进行一些调整，见构建TensorRT引擎)；将model.onnx拷贝到Jeston Nano设备上，可以通过Jetpack系统自带的trtexec命令行工具，或者自己编写构建程序，将model.onnx转换并导出为TensorRT引擎文件model.engine；推理程序加载model.engine和输入数据，进行模型推理