C++ API 文档#

注意

这是 NVIDIA TensorRT 库的 TensorRT C++ API。NVIDIA TensorRT C++ API 允许开发人员使用 C++ 导入、校准、生成和部署网络。网络可以直接从 ONNX 导入。它们也可以通过实例化各个层并直接设置参数和权重以编程方式创建。

本节说明了 C++ API 的基本用法,假设您从 ONNX 模型开始。sampleOnnxMNIST 示例更详细地说明了此用例。

可以通过头文件 NvInfer.h 访问 C++ API,它位于 nvinfer1 命名空间中。例如,一个简单的应用程序可能以以下内容开始

#include “NvInfer.h”

using namespace nvinfer1;

TensorRT C++ API 中的接口类以 I 前缀开头,例如 ILoggerIBuilder

如果之前不存在 CUDA 上下文,则在 TensorRT 首次调用 CUDA 时会自动创建 CUDA 上下文。但是,通常最好在首次调用 TensorRT 之前创建和配置您自己的 CUDA 上下文。

本章中的代码未使用智能指针来说明对象生命周期;但是,建议将它们与 TensorRT 接口一起使用。

构建阶段#

要创建构建器,您首先必须实例化 ILogger 接口。此示例捕获所有警告消息,但忽略信息性消息

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

然后,您可以创建构建器的实例

IBuilder* builder = createInferBuilder(logger);

创建网络定义#

创建构建器后,优化模型的第一步是创建网络定义。网络创建选项是使用标志组合 OR 在一起指定的。

您可以使用 NetworkDefinitionCreationFlag::kSTRONGLY_TYPED 标志指定网络应被视为强类型。有关更多信息,请参阅 强类型网络 部分。

最后,创建网络

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

从头开始创建网络定义(高级)#

除了使用解析器之外,您还可以通过网络定义 API 直接向 TensorRT 定义网络。此方案假定每层权重已在主机内存中准备就绪,以便在网络创建期间传递给 TensorRT。

此示例创建一个简单的网络,其中包含输入、卷积、池化、矩阵乘法、Shuffle、激活和 Softmax 层。它还将权重加载到 weightMap 数据结构中,该结构在以下代码中使用。

首先,创建构建器和网络对象。请注意,记录器是使用 logger.cpp 文件初始化的,该文件是所有 C++ 示例通用的。C++ 示例辅助类和函数可以在 common.h 头文件中找到。

auto builder = SampleUniquePtr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(sample::gLogger.getTRTLogger()));
auto network = SampleUniquePtr<nvinfer1::INetworkDefinition>(0);

通过指定输入张量的名称、数据类型和完整维度,将输入层添加到网络。网络可以有多个输入,尽管在此示例中只有一个

auto data = network->addInput(INPUT_BLOB_NAME, datatype, Dims4{1, 1, INPUT_H, INPUT_W});

添加具有隐藏层输入节点、步幅以及滤波器和偏差权重的卷积层。

auto conv1 = network->addConvolution(
*data->getOutput(0), 20, DimsHW{5, 5}, weightMap["conv1filter"], weightMap["conv1bias"]);
conv1->setStride(DimsHW{1, 1});

注意

传递给 TensorRT 层的权重位于主机内存中。

添加池化层;请注意,前一层的输出作为输入传递。

auto pool1 = network->addPooling(*conv1->getOutput(0), PoolingType::kMAX, DimsHW{2, 2});
pool1->setStride(DimsHW{2, 2});

添加 Shuffle 层以重塑输入,为矩阵乘法做准备

int32_t const batch = input->getDimensions().d[0];
int32_t const mmInputs = input.getDimensions().d[1] * input.getDimensions().d[2] * input.getDimensions().d[3];
auto inputReshape = network->addShuffle(*input);
inputReshape->setReshapeDimensions(Dims{2, {batch, mmInputs}});

现在,添加一个矩阵乘法层。模型导出器提供了转置权重,因此指定了 kTRANSPOSE 选项。

IConstantLayer* filterConst = network->addConstant(Dims{2, {nbOutputs, mmInputs}}, mWeightMap["ip1filter"]);
auto mm = network->addMatrixMultiply(*inputReshape->getOutput(0), MatrixOperation::kNONE, *filterConst->getOutput(0), MatrixOperation::kTRANSPOSE);

添加偏差,它将在批次维度上广播。

auto biasConst = network->addConstant(Dims{2, {1, nbOutputs}}, mWeightMap["ip1bias"]);
auto biasAdd = network->addElementWise(*mm->getOutput(0), *biasConst->getOutput(0), ElementWiseOperation::kSUM);

添加 ReLU 激活层

auto relu1 = network->addActivation(*ip1->getOutput(0), ActivationType::kRELU);

添加 SoftMax 层以计算最终概率

auto prob = network->addSoftMax(*relu1->getOutput(0));

为 SoftMax 层的输出添加一个名称,以便可以在推理时将张量绑定到内存缓冲区

prob->getOutput(0)->setName(OUTPUT_BLOB_NAME);

将其标记为整个网络的输出

network->markOutput(*prob->getOutput(0));

现在已经完全构建了表示 MNIST 模型的网络。有关如何构建引擎并使用此网络运行推理的说明,请参阅 构建引擎执行推理 部分。

有关层的更多信息,请参阅 TensorRT 算子文档

使用 ONNX 解析器导入模型#

现在,必须从 ONNX 表示形式填充网络定义。ONNX 解析器 API 在文件 NvOnnxParser.h 中,解析器位于 nvonnxparser C++ 命名空间中。

#include “NvOnnxParser.h”

using namespace nvonnxparser;

您可以创建一个 ONNX 解析器来填充网络,如下所示

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

然后,读取模型文件并处理任何错误。

parser->parseFromFile(modelFile,
    static_cast<int32_t>(ILogger::Severity::kWARNING));
for (int32_t i = 0; i < parser->getNbErrors(); ++i)
{
std::cout << parser->getError(i)->desc() << std::endl;
}

TensorRT 网络定义的一个重要方面是它包含指向模型权重的指针,构建器会将这些指针复制到优化的引擎中。由于网络是使用解析器创建的,因此解析器拥有权重占用的内存,因此在构建器运行之后才能删除解析器对象。

构建引擎#

下一步是创建构建配置,指定 TensorRT 应如何优化模型。

IBuilderConfig* config = builder->createBuilderConfig();

此接口具有许多属性,您可以设置这些属性来控制 TensorRT 如何优化网络。一个重要的属性是最大工作区大小。层实现通常需要临时工作区,此参数限制了网络中任何层可以使用的最大大小。如果提供的工作区不足,TensorRT 可能无法为层找到实现。默认情况下,工作区设置为给定设备的全局内存总大小;必要时限制它,例如,当要在单个设备上构建多个引擎时。

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

另一个重要的考虑因素是 CUDA 后端实现的最大共享内存分配。当 TensorRT 需要与其他应用程序共存时,例如当 TensorRT 和 DirectX 同时使用 GPU 时,此分配至关重要。

config->setMemoryPoolLimit(MemoryPoolType::kTACTIC_SHARED_MEMORY, 48 << 10);

指定配置后,即可构建引擎。

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

由于序列化引擎包含权重的必要副本,因此不再需要解析器、网络定义、构建器配置和构建器,可以安全地删除它们

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

然后,可以将引擎保存到磁盘,并且可以删除序列化到其中的缓冲区。

delete serializedModel

注意

序列化引擎不跨平台移植。它们特定于构建它们的精确 GPU 型号(以及平台)。

构建引擎旨在作为离线过程,因此可能需要相当长的时间。优化构建器性能 部分提供了有关加快构建器运行速度的提示。

反序列化计划#

当您有一个先前序列化的优化模型并想要执行推理时,您必须首先创建 Runtime 接口的实例。与构建器一样,运行时需要记录器的实例

IRuntime* runtime = createInferRuntime(logger);

TensorRT 提供了三种主要的反序列化引擎的方法,每种方法都有其用例和优点。

内存中反序列化

此方法简单明了,适用于较小的模型或当内存不是约束时。

std::vector<char> modelData = readModelFromFile("model.plan");
ICudaEngine* engine = runtime->deserializeCudaEngine(modelData.data(), modelData.size());

IStreamReaderV2 反序列化

此方法允许更受控地读取引擎文件,并且对于自定义文件处理或权重流式传输非常有用。它支持从主机和设备指针读取,并可能实现性能改进。使用此方法,无需将整个计划文件读取到要反序列化的缓冲区中,因为 IStreamReaderV2 允许根据需要分块读取文件,并可能绕过 CPU,从而降低峰值 CPU 内存使用率。

class MyStreamReaderV2 : public IStreamReaderV2 {
    // Custom implementation with support for device memory reading
};
MyStreamReaderV2 readerV2("model.plan");
ICudaEngine* engine = runtime->deserializeCudaEngine(readerV2);

IStreamReaderV2 方法对于大型模型或使用 GPUDirect 或权重流式传输等高级功能时尤其有利。它可以显着减少引擎加载时间和内存使用量。

选择反序列化方法时,请考虑您的具体要求

  • 对于小型模型或简单的用例,内存中反序列化通常就足够了。

  • 对于大型模型或内存效率至关重要时,请考虑使用 IStreamReaderV2

  • 如果您需要自定义文件处理或权重流式传输功能,IStreamReaderV2 提供了必要的灵活性。

执行推理#

引擎保存优化的模型,但您必须管理中间激活的其他状态才能执行推理。这是使用 ExecutionContext 接口完成的

IExecutionContext *context = engine->createExecutionContext();

一个引擎可以有多个执行上下文,允许一组权重用于多个重叠的推理任务。(当前的一个例外是使用动态形状时,除非指定了预览功能 kPROFILE_SHARING_0806,否则每个优化配置文件只能有一个执行上下文。)

要执行推理,您必须为输入和输出传递 TensorRT 缓冲区,TensorRT 要求您使用对 setTensorAddress 的调用来指定这些缓冲区,该调用接受张量的名称和缓冲区的地址。您可以使用您为输入和输出张量提供的名称查询引擎,以找到数组中的正确位置

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

如果引擎是使用动态形状构建的,您还必须指定输入形状

context->setInputShape(INPUT_NAME, inputDims);

然后,您可以调用 TensorRT 的方法 enqueueV3 以使用 CUDA 流启动推理

context->enqueueV3(stream);

网络将异步执行或不异步执行,具体取决于网络的结构和功能。可能导致同步行为的非详尽功能列表包括数据相关的形状、DLA 使用、循环和同步插件。通常在内核之前和之后使用 cudaMemcpyAsync() 排队数据传输,以便在数据尚未存在于 GPU 中时将其从 GPU 移出。

要确定内核(以及可能的 cudaMemcpyAsync())何时完成,请使用标准的 CUDA 同步机制,例如事件或等待流。