Windows上运行TensorRT的“Hello World“代码

本文将分别介绍C++版本和Python版本。

一、C++版本

代码地址：“Hello World” For TensorRT From ONNX

这个项目使用TensorRT将一个已经训练好的ONNX模型转换为trt文件，并进行推理。模型在MNIST - Handwritten Digit Recognition数据集上进行训练。

总的来说，该项目分为三个步骤：

• 将ONNX模型转为TensorRT网络
• 创建engine
• 进行推理

1.将ONNX模型转为TensorRT网络

这一步需要创建一个ONNX解析器（parser），然后传入ONNX模型地址，将其转换为TensorRT网络。

首先初始化parser，将目标TensorRT网络和记录器对象（logger）作为输入。

auto parser = nvonnxparser::createParser(*network, sample::gLogger.getTRTLogger());

接着将ONNX文件传入parser，如果解析失败，则打印相应级别的日志。

if (!parser->parseFromFile(model_file, static_cast<int>(sample::gLogger.getReportableSeverity())))
{
	  string msg("failed to parse onnx file");
	  sample::gLogger->log(nvinfer1::ILogger::Severity::kERROR, msg.c_str());
	  exit(EXIT_FAILURE);
}

如果解析成功的话，TensorRT网络被构建了出来。下面就可以创建TensorRT引擎来执行推理了。

2.创建engine

首先需要创建生成器（build），需要一个记录器作为输入，用于报告网络中的错误、警告和信息性消息：

IBuilder* builder = createInferBuilder(sample::gLogger);

接着从TensorRT网络中生成引擎engine：

SampleUniquePtr<IHostMemory> plan{builder->buildSerializedNetwork(*network, *config)};

然后就可以通过engine进行推理了。

3.进行推理

虽然我们使用engine保存了一个优化后的模型，但是在实际推理的过程中，还需要额外保存每一层的输出结果，因此我们首先需要创建一个执行上下文（context）：

IExecutionContext *context = engine->createExecutionContext();

一个engine可以有多个context，从而允许一组权重被多个重叠的推理过程共用（使用dynamic shapes是一个例外）。

执行推断前，需要明确输入输出的张量名称和缓冲区地址：

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

然后就可以调用TensorRT的enqueueV3方法使用CUDA流来启动推断：

context->enqueueV3(stream);

4.如何异步

网络是否异步执行取决于网络的结构与特性。

导致同步的情况包括数据依赖形状（模型的某些操作或层的输出形状会根据输入数据的实际形状变化）、DLA使用（深度学习工作负载设计的硬件加速器，内部有自己的同步机制）、循环（loops）和插件（plugins）等。

通常我们使用cudaMemcpyAsync()方法将数据异步输入和输出GPU，通过GPU异步计算。

二、Python版本

代码地址：Introduction To Importing ONNX Models Into TensorRT Using Python

首先导入tensorrt api包：

import tensorrt as trt

1.构建阶段 

创建日志logger：

logger = trt.Logger(trt.Logger.WARNING)

创建构造器builder：

builder = trt.Builder(logger)

2.网络定义

network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

3.网络导出

创建ONNX解析器用以填充网络：

parser = trt.OnnxParser(network, logger)

然后解析网络，并解决产生的错误信息：

success = parser.parse_from_file(model_path)
for idx in range(parser.num_errors):
    print(parser.get_error(idx))

if not success:
    pass # Error handling code here

4.创建engine

首先创建一个构建配置（config），指定TensorRT应如何优化模型：

config = builder.create_builder_config()

该接口有许多属性，例如指定工作空间的大小。默认情况下，engine的工作空间是给定设备全局内存大小，如果有多个engine在一个设备上构建时，有必要限制它，例如我们可以将其工作空间限制为1MB：

config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 20) # 1 MiB

接着就可以序列化生成engine：

serialized_engine = builder.build_serialized_network(network, config)