TensorRT部署YOLOv5-03-TensorRT

发表于 | 分类于 ， |

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

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

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

构建一个推理引擎有3种方式，这些方式各有优劣

  • 框架自带TRT接口

    英伟达与部分深度学习框架有合作，例如Tensorflow内置了TF-TRT，Pytorch有Torch-TensorRT，使用这些框架内置的TensorRT接口，可以将训练好的模型无缝衔接，直接在原有框架上调用对应接口进行推理，这种方式的优势在于非常方便，环境统一，开发效率较高，遇到不支持的算子会返回到原框架的实现。缺点是性能较差，不能最大限度的利用TensorRT的优化加速能力，另外在很多资源有限的嵌入式设备上安装Tensorflow这种较大的深度学习工具也不现实，因此这种方式通常是应用于服务端推理，本文不对这种方式进行介绍

  • 使用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搭建网络等方面不进行介绍