TensorRT基本流程概览

1、核心API

[C++ API](API Reference :: NVIDIA Deep Learning TensorRT Documentation)

TensorRT C++ API允许开发者导入、校准、生成和部署网络。网络可以直接从ONNX导入,或者用C++实例化每一个层构建网络结构,然后将参数和权重(储存在文件中)装载到结构中。

核心API主要分为以下一个部分,这些API都被包含在NvInfer.h头文件中:

(关于C++ API及示例代码等更多信息,参考Developer Guide :: NVIDIA Deep Learning TensorRT Documentation

2、TensorRT流程实例

现在我们有一个onnx模型文件(pytorch、tensorflow等都提供了转onnx的工具),如何使用TensorRT进行推理加速呢?首先我们需要知道该过程大体上分为两步:

  • 构建engine文件(即NEURAL ETWORK到PLAN这部分)。TensorRT(TRT)能进行推理加速的原因之一时将模型针对具体的硬件进行了优化构建的,所以TRT首先需要将ONNX转为可在GPU上高效执行的格式,并且在此过程中进行优化。这个过程的最终结构就是得到一个engine格式的序列化文件。但是这个转换过程比较耗时,我们不希望每次启动都需要长时间等待它转换,所以我们可以将转换的结果以二进制方式保存到本地文件,通常是以.engine结尾的文件。
  • 使用engine进行推理(即从PLAN到 Validate using TensorRT这部分)。启动时直接从本地文件读取二进制的engine文件,再反序列化就可以进行相关的推理任务。

下面将分别讲解两个部分,内容基本来自The C++ API

2.1、构建engine文件

C++ API在头文件NvInfer.h中,并且在nvinfer1命名空间中,所以当使用C++的TensorRT时,有如下开头:

#include "NvInfer.h"

using namespace nvinfer1;

TensorRT的API中,接口类是以I为前缀的,比如ILoggerIBuilder等等。

如果在 TensorRT 首次调用 CUDA 之前没有 CUDA context,则会自动创建 CUDA context。一般情况下,最好在首次调用 TensorRT 之前手动创建和配置 CUDA context。

2.1.1、创建构建器

构建器顾名思义用来构建的,比如从构建器构建网络定义

我们需要创建一个构建器,但在这之前必须要实例化一个ILogger接口(通过继承ILogger并实现其中的log函数),比如下面的示例会捕获所有警告信息,但是忽略informational messages(系统状态、操作进展或其他非错误性事件的信息)。然后在实例化构建器时将log对象传入。即创建构建器必须要有ILogger对象:

class Logger:pubilc ILogger{
    void log(Severity severity, const char*msg) noexcept override{
        // 阻止info-level级别信息
        if(severity<=Severity::kWARNING){
            std::cout<<msg<<std::endl;
        }
    }
} logger;

// 创建构建器
IBuilder* builder = createInferBuilder(logger);

上面的代码中,severity指示日志消息的严重程度,msg是日志消息的内容,noexcept表示该函数不抛出异常,override表明是对同名函数的重写。

2.1.2、用构建器创建一个网络定义

创建构建器后,第一步是创建网络定义,如:

uint32_t flag = 1U<<static_cast<uint32_t>(NetworkDefinitionCreateionFlag::KEXPLICIT_BATCH);

INetworkDefinition* network =builder->createNetworkV2(flag);

或者:

INetworkDefinition* network =builder->createNetworkV2(1U <<
        static_cast<int>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH));

如果要使用ONNX parser必须要使用kEXPLICIT_BATCH标志,详情见:6.9. Explicit Versus Implicit Batch

在RTX中,有两个阶段:1)build engine:将储存了层名和权重的wts文件转换为engine文件。该过程需要IBuilder和INetwork;2)inference:直接将engine读取到IRuntime中,进行推理。

此时的网络定义只是一个空壳,没有结构也没有权重参数。

2.1.3、使用ONNX Parser向网络定义中迁移结构和参数

ONNX是一种中间件表示协议,规定了网络的结构且储存了参数。现在可以用TensorRT的ONNX Parser接口将ONNX文件中的结构和参数填充到上文定义的空壳网络network中。

首先ONNX parser的API位于NvOnnxParser.h头文件中,parser位于nvonnxparser命名空间:

#include "NvOnnxParser.h"

using namespace nvonnxparser;

要使用parser进行填充,首先需要创建一个parser实例:

IParser* parser=createParser(*network,logger);

显然任何需要日志的接口在实例化时都需要ILogger对象。

现在我们读取ONNX模型文件并且处理出现的任何错误:

parser->parserFromFile(modelFile, static_cast<int32_t>(ILogger::Serverity::kWARNING));

for(int32_t i=0;i<parser.getNbErrors();++i){
    std::cout<<parser->getError(i)->desc()<<std::endl;
}

在这里我们尝试读取本地的ONNX文件,同时定义了需要捕获的错误的等级。函数在读取文件发送的所有小于该等级的错误都会被记录。后续通过getNbErrors和getError获取到错误的详细描述,帮助我们了解模型的读取情况。

2.1.4、构建一个engine文件

TensorRT能够有效加速的一个原因是,会根据硬件的具体情况,将ONNX等其他模型进行构建优化适配。但是每一次构建都很耗时,所以构建完毕后通常都是序列化保存到本地文件,文件以.engine为后缀。下一次推理初始化时只需要反序列化加载engine文件就ok了,不需要再进行耗时的优化编译了。

需要说明的时,engine是TensorRT根据你的硬件、CUDA及cuDNN针对性的优化,所以本机构建的engine文件只能在本机使用。就算是两台软硬件相同的设备也无法混用(没有具体测试过,但是两台相同GPU的设备测试过确实会报错)。甚至同一个设备两次构建engine也不一定相同。

在上节中我们读取了ONNX模型文件,现在我们尝试将其构建为设配设备的engine文件。

首先在构建时,我们需要使用IBuilderConfig设定一些构建优化的细节:

IBuilderCOnfig* config=builder->createBuilderConfig();

config->setMemoryPoolLimit(MemoryPoolType::kWORKSPACE, 1U<20);

IBuilderConfig API

最后就可以构建engine了:

IHostMemory* serializedModel = builder->buildSerializedNetwork(*network, *config);

由于序列化的engine包含了包含了必需的权重,而parser、network definition、builder config和builder不再需要了,因此可以安全删除:

delete parser;
delete network;
delete config;
delete builder;

2.1.5、保存序列化engine文件

上节中我们构建了序列化的engine保存在IHostMemory中,构建序列化文件时很耗时,我们并不希望下次启动时还要序列化一次。所以我们可以将engine以文件的形式保存到本地(已经针对设备优化过的用于TensorRT推理,和ONNX文件不同)。

保存这部分实际上和TensorRT没有什么关系了,就是借助C++的文件操作进行保存:

// Serialize the engine
IHostMemory* serialized_engine = engine->serialize();
assert(serialized_engine != nullptr);

// Save engine to file
std::ofstream p("saved.engine", std::ios::binary);
if (!p) {
std::cerr << "Could not open plan output file" << std::endl;
assert(false);
}
p.write(reinterpret_cast<const char*>(serialized_engine->data()), serialized_engine->size());

delete serializedModel  // 最后安全删除

读取engine文件的代码如下,其中engine和runtime需要在TensorRT推理部分介绍:

std::ifstream file(engine_name, std::ios::binary);
if (!file.good()) {
std::cerr << "read " << engine_name << " error!" << std::endl;
assert(false);
}
size_t size = 0;
file.seekg(0, file.end);
size = file.tellg();
file.seekg(0, file.beg);
char* serialized_engine = new char[size];
assert(serialized_engine);
file.read(serialized_engine, size);
file.close();


IRuntime* runtime = createInferRuntime(logger);
ICudaEngine* engine = runtime->deserializeCudaEngine(serialized_engine, size);

2.2、使用engine执行推理

TODO:3.3. Performing Inference

在推理时,我们需要一个运行环境进行相关的控制。所以必须先实例化一个Runtime接口,同样的它也需要一个ILogger

IRuntime* runtime=createInferRuntime(logger);

通过IRuntime我们可以创建一个ICudaEngine,ICudaEngine有优化后的模型,保存的是网络结构和参数,所以在创建时需要给定engine文件数据加载进数据进行创建:

ICudaEngine* engine=runtime->deserializeCudaEngine(serialized_engine,size);

在推理时会产生很多中间值,而ICudaEngine只是持有模型数据,不具有管理中间值和运行状态的能力,而这部分的功能是由IExecutionContext的执行上下文实现的:

IExecutionContext *context = engine->createExecutionContext();

为了高效的数据传递,我们需要给模型的输入输出节点绑定一个buffer,推理时输入直接从绑定的输入buffer读取数据,输出结果也直接给绑定的输出buffer,提高推理效率:

context->setTensorAddress(INPUT_NAME, inputBuffer);
context->SetTensorAddress(OUTPUT_NAME, outputBuffer);

其中INPUT_NAME和OUTPUT_NAME是模型输入和输出节点的名字,即通过名字定位到输入输出节点,然后绑定buffer。

在推理之前我们需要将输入数据填充到inputBuffer中,然后调用 TensorRT 的 enqueueV3 方法,开始使用 CUDA 流进行推理::

context->enqueueV3(stream);

转载请注明来源,欢迎对文章中的引用来源进行考证,欢迎指出任何有错误或不够清晰的表达。可以在下面评论区评论,也可以邮件至 ishyj@qq.com