什么是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;
流程图
总结
本文介绍了使用 TensorRT 的必备接口,掌握了这些接口之间的调用关系,也就明白了 TensorRT 的工作流程,后续将结合实际工程,详细介绍每一步骤的具体细节。
参考
https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html