TensorRT原理与核心接口-优快云博客

什么是TensorRT

TensorRT是由Nvidia推出的C++语言开发的高性能神经网络推理库，是一个用于生产部署的优化器和运行时引擎。其高性能计算能力依赖于Nvidia的图形处理单元。它专注于推理任务，与常用的神经网络学习框架形成互补，包括TensorFlow、Caffe、PyTorch、MXNet等。可以直接载入这些框架的已训练模型文件，也提供了API接口通过编程自行构建模型。

TensorRT依赖于Nvidia的深度学习硬件环境，可以是GPU也可以是DLA，如果没有的话则无法使用。

TensorRT支持目前大部分的神经网络Layer的定义，同时提供了API让开发者自己实现特殊Layer的操作，相关函数为 INetworkDefinition::addPluginV2()。

关键接口类型

TensorRT核心库中，最关键的几种接口类型有：

IExecutionContext 推理引擎运行上下文
ICudaEngine 推理引擎
IRuntime CudaEngine反序列化
INetWorkDefinition 网络定义
IParser 网络模型解析
IOptimizationProfile 优化配置
IBuilderConfig CudaEngine的构造参数
IBuilder 构造器，主要用于构造CudaEngine
ILogger 日志接口，需要开发者实现

IExecutionContext

推理引擎运行上下文(Context)，使用CudaEngine进行推理操作，是推理操作的最终执行接口。

允许一个CudaEngine具有多个运行Context，每个Context可以使用不同的Batch大小，如果网络模型的输入尺寸支持动态调整，那么每个上下文还可以使用各自不同的尺寸输入。

主要功能为执行推理操作，具体函数 IExecutionContext::executeV2(bindings)

通过CudaEngine创建，ICudaEngine::createExecutionContext()

ICudaEngine

可称为推理引擎，允许应用程序调用这个接口来执行推理，支持同步执行和异步执行，通过Cuda中的Stream和Event来实现异步。一个推理引擎可以有多个运行Context，并且支持批量输入执行。

CudaEngine的主要作用就是通过创建Context执行推理任务，创建Context的方式为 ICudaEngine::createExecutionContext()。

CudaEngine可以序列化到内存中，然后缓存到磁盘上，下次使用的时候可以直接从磁盘文件载入到内存再序列为CudaEngine即可，可以省去很多时间和参数配置等一下操作。

CudaEngine的创建方式依赖于 INetwrorkDefinition(网络定义接口)，NetWorkDefinition 一般通过解析ONNX模型文件或TensorFlow的已训练模型文件获得。ONNX模型文件的解析需要用到 nvonnxparser::IParser 接口。

注意：从 NeworkDefinition 生成 CudaEngine 是一个很耗时的过程，可以将生成的 CudaEngine 缓存到磁盘文件，后续使用的时候直接载入即可使用。

另外，CudaEngine 是不可以跨平台的，不同的 GPU 型号和不同的 TensorRT 版本生成的 CudaEngine 可能都会不兼容，在使用缓存的 CudaEngine 文件的时候要注意区分，可将这些影响因素作为该文件名的一部分进行区分，例如 Win64_RTX2080TI_70011.engine 这类的文件名。

相关接口：

IExecutionContext，生成 Context 并通过Context进行推理，相关函数：ICudaEngine::createExecutionContext()
IRuntime，反序列化缓存文件，得到CudaEngine，相关函数：IRuntime::deserializeCudaEngine()
IBuilder，根据 NetworkDefinision 和 BuilderConfig 构造 CudaEngine，相关函数：IBuilder::buildEngineWithConfig(INetworkDefinision,IBuilderConfig)

IRuntime

这个接口的名字起得容易让人误解，看名字好像很底层的一个接口，但它的实际作用，主要只有一个，就是将 CudaEngine 的序列化缓存文件反序列化回来，重新得到 CudaEngine 对象。

获取方式：nvinfer1::createInferRuntime(void)

相关接口：

ICudaEngine，根据 NetworkDefinision 和 BuilderConfig 构造 CudaEngine，相关函数：IBuilder::buildEngineWithConfig()

INetWorkDefinition

网络定义接口，这个接口提供了一些列的函数让开发者可以从头构造一个神经网络，包括输入输出张量的维度大小，每个Layer的类型和激活函数等等，还提供了接口让开发者添加自定义的Layer，是一个很强大的接口。但是，在一般的实际使用过程中基本上不会用到它的任何一个函数。因为网络的定义是根据已训练网络模型文件例如ONNX文件自动生成的。

该接口的一般使用步骤：

通过 IBuilder::createNetwork() 生成一个 INetWorkDefinition。
使用接口 NvOnnxParser::createParser(&INetWorkDefinition,...) 创建一个与该 INetWorkDefinition 绑定的 IPaser 对象。
调用 IPaser::parseFromFile("path.onnx")，将根据模型文件构造 INetWorkDefinition 对象。

IParser

ONNX Parser，解析已训练模型文件，根据模型文件构造绑定的 INetWorkDefinition 对象。

获取方式：NvOnnxParser::createParser(&INetWorkDefinition,...)

主要函数：IPaser::parseFromFile("path.onnx")

IOptimizationProfile

为动态模型指定每一个输入输出张量的维度，函数为 IOptimizationProfile::setDimensions() 。

在构造 CudaEngine 的时候至少要有一个 IOptimizationProfile，因为每个 ExecutionContext 在使用之前都要先指定一个 IOptimizationProfile 才可以执行推理操作。

获取方式 IBuilder::createOptimizationProfile(void)

相关接口：

IBuilderConfig，每个 IBuilderConfig 至少要有一个 IOptimizationProfile，IOptimizationProfile 将随着 IBuilderConfig 被构造到 CudaEngine 中，被 ExecutionContext 指定使用。相关函数为 IBuilderConfig::addOptimizationProfile(IOptimizationProfile)
IExecutionContext，每个 Context 被创建之后需要先指定一个 IOptimizationProfile，IExecutionContext::setOptimizationProfile(index)，这个 index 是 IOptimizationProfile 在 IBuilderConfig 中的序号，顺序是按照 addOptimizationProfile 的调用次序来的。

IBuilderConfig

构造 CudaEngine 的配置参数，可添加 IOptimizationProfile 配置，设置最大工作内存空间、最大Batch大小、最小可接受精度级别、半浮点精度运算等。

获取方式 IBuilder::createBuilderConfig(void)

相关接口：

IBuilder，根据 NetworkDefinision 和 BuilderConfig 构造 CudaEngine，函数：IBuilder::buildEngineWithConfig(INetworkDefinision,IBuilderConfig)

IBuilder

IBuilder 接口主要用来构建 CudaEngine，也用来生成 INetWorkDefinition 接口对象、IOptimizationProfile 和 IBuilderConfig 接口对象。

ILogger

日志接口，用来输出 TensorRT 内部的一些消息、警告、错误等信息。

在创建 IBuilder 和 IRuntime 的时候都会需要传入 ILogger 对象，我们需要实现这个接口并创建一个对象传给它们。最简单的一个接口实现如下：

class Logger : public ILogger           
 {
     void log(Severity severity, const char* msg) override
     {
         // suppress info-level messages
         if (severity != Severity::kINFO)
             std::cout << msg << std::endl;
     }
 } gLogger;