TensorRT 的 C++ API 使用详解

1. TensorRT 的 C++ API 使用示例

进行推理，需要先创建IExecutionContext对象，要创建这个对象，就需要先创建一个ICudaEngine的对象（engine）。

两种创建engine的方式：

使用模型文件创建engine，并可把创建的engine序列化后存储到硬盘以便后面直接使用；
使用之前已经序列化存储的engine，这种方式比较高效些，因为解析模型并生成engine还是挺慢的。

无论哪种方式，都需要创建一个全局的iLogger对象，并被用来作为很多TensorRT API方法的参数使用。如下是一个logger创建示例：

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;

2. 用 C++ API 创建TensorRT网络

2.1. 使用 C++ 的 parser API 导入模型

1. 创建TensorRT builder和network

IBuilder* builder = createInferBuilder(gLogger);
nvinfer1::INetworkDefinition* network = builder->createNetwork();

2.针对特定格式创建TensorRT parser

// ONNX
auto parser = nvonnxparser::createParser(*network,
        gLogger);
// UFF
auto parser = createUffParser();
// NVCaffe
ICaffeParser* parser = createCaffeParser();

3. 使用parser解析导入的模型并填充network

parser->parse(args);

具体的args要看使用什么格式的parser。

必须在网络之前创建构建器，因为它充当网络的工厂。不同的解析器在标记网络输出时有不同的机制。

2.2. 使用 C++ Parser API 导入 Caffe 模型

1. 创建builder和network

IBuilder* builder = createInferBuilder(gLogger);
INetworkDefinition* network = builder->createNetwork();

2. 创建caffe parser

ICaffeParser* parser = createCaffeParser();

3. 解析导入的模型

const IBlobNameToTensor* blobNameToTensor = parser->parse("deploy_file", 
              "modelFile", 
              *network, 
              DataType::kFLOAT);

这将把Caffe模型填充到TensorRT的network。最后一个参数指示解析器生成权重为32位浮点数的网络。使用DataType :: kHALF将生成具有16位权重的模型。

除了填充网络定义之外，parser还返回一个字典，该字典是从 Caffe 的 blob names 到 TensorRT 的 tensors 的映射。与Caffe不同，TensorRT网络定义没有in-place的概念。当Caffe模型使用in-place操作时，字典中返回的相应的TensorRT tensors是对那个blob的最后一次写入。例如，如果是一个卷积写入到了blob并且后面跟的是ReLU，则该blob的名字映射到TensorRT tensors就是ReLU的输出。

4. 给network分配输出

for (auto& s : outputs)
    network->markOutput(*blobNameToTensor->find(s.c_str()));

2.3. 使用 C++ Parser API 导入 TensorFlow 模型

对于一个新的工程，推荐使用集成的TensorFlow-TensorRT作为转换TensorFlow network到TensorRT的方法来进行推理。具体可参考：Integrating TensorFlow With TensorRT

从TensorFlow框架导入，需要先将TensorFlow模型转换为中间格式：UFF（Universal Framework Format）。相关转换可参考Coverting A Frozen Graph to UFF

更多关于UFF导入的信息可以参考：https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html#mnist_uff_sample

先来看下如何用C++ Parser API来导入TensorFlow模型。

1. 创建builder和network

IBuilder* builder = createInferBuilder(gLogger);
INetworkDefinition* network = builder->createNetwork();

2. 创建UFF parser

IUFFParser* parser = createUffParser();

3. 向UFF parser声明network的输入和输出

parser->registerInput("Input_0", DimsCHW(1, 28, 28), UffInputOrder::kNCHW);
parser->registerOutput("Binary_3");

注意：TensorRT期望的输入tensor是CHW顺序的。从TensorFlow导入时务必确保这一点，如不是CHW，那就先转换成CHW顺序的。

4. 解析已导入的模型到network

parser->parse(uffFile, *network, nvinfer1::DataType::kFLOAT);

2.4. 使用 C++ Parser API 导入 ONNX 模型

使用限制：注意版本问题，TensorRT5.1 附带的 ONNX Parser支持的ONNX IR版本是0.0.3，opset版本是9。通常，较新的ONNX parser是后向兼容的。更多信息可参考：ONNX Model Opset Version Converter 和 onnx-tensorrt。

更多关于ONNX导入的信息也可参考：
https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html#onnx_mnist_sample

1. 创建ONNX parser，parser使用辅助配置管理SampleConfig对象将输入参数从示例的可执行文件传递到parser对象

nvonnxparser::IOnnxConfig* config = nvonnxparser::createONNXConfig();
//Create Parser
nvonnxparser::IONNXParser* parser = nvonnxparser::createONNXParser(*config);

2. 填充模型

parser->parse(onnx_filename, DataType::kFLOAT);

3. 转换模型到TensorRT的network

parser->convertToTRTNetwork();

4. 从模型获取network

nvinfer1::INetworkDefinition* trtNetwork = parser->getTRTNetwork();

3. 用 C++ API 构建 engine

下一步是调用TensorRT的builder来创建优化的runtime。 builder的其中一个功能是搜索其CUDA内核目录以获得最快的实现，因此用来构建优化的engine的GPU设备和实际跑的GPU设备一定要是相同的才行。

builder具有许多属性，可以通过设置这些属性来控制网络运行的精度，以及自动调整参数。还可以查询builder以找出硬件本身支持的降低的精度类型。

有两个特别重要的属性：最大batch size和最大workspace size。

最大batch size指定TensorRT将要优化的batch大小。在运行时，只能选择比这个值小的batch。
各种layer算法通常需要临时工作空间。这个参数限制了网络中所有的层可以使用的最大的workspace空间大小。如果分配的空间不足，TensorRT可能无法找到给定层的实现。

1. 用builder对象创建构建engine

builder->setMaxBatchSize(maxBatchSize);
builder->setMaxWorkspaceSize(1 << 20);
ICudaEngine* engine = builder->buildCudaEngine(*network);

2. 用完分配过的network，builder和parser记得解析

parser->destroy();
network->destroy();
builder->destroy();

4. 用 C++ API 序列化一个模型

序列化模型，即把engine转换为可存储的格式以备后用。推理时，再简单的反序列化一下这个engine即可直接用来做推理。通常创建一个engine还是比较花时间的，可以使用这种序列化的方法避免每次重新创建engine。

注意：序列化的engine不能跨平台或在不同版本的TensorRT间移植使用。因为其生成是根据特定版本的TensorRT和GPU的。

1. 序列化

IHostMemory *serializedModel = engine->serialize();
// store model to disk
// <…>
serializedModel->destroy();

2. 创建一个runtime并用来反序列化

IRuntime* runtime = createInferRuntime(gLogger);
ICudaEngine* engine = runtime->deserializeCudaEngine(modelData, modelSize, nullptr);

5. 用 C++ API 执行推理

1. 创建一个Context用来存储中间激活值

IExecutionContext *context = engine->createExecutionContext();

一个engine可以有多个execution context，并允许将同一套weights用于多个推理任务。可以在并行的CUDA streams流中按每个stream流一个engine和一个context来处理图像。每个context在engine相同的GPU上创建。

2. 用input和output的blob名字获取对应的input和output的index

int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);
int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);

3. 使用上面的indices，在GPU上创建一个指向input和output缓冲区的buffer数组

void* buffers[2];
buffers[inputIndex] = inputbuffer;
buffers[outputIndex] = outputBuffer;

4. 通常TensorRT的执行是异步的，因此将kernels加入队列放在CUDA stream流上

context.enqueue(batchSize, buffers, stream, nullptr);

通常在kernels之前和之后来enquque异步memcpy()以从GPU移动数据（如果尚未存在）。

enqueue()的最后一个参数是一个可选的CUDA事件，当输入缓冲区被消耗且它们的内存可以安全地重用时这个事件便会被信号触发。

为了确定kernels（或可能存在的memcpy()）何时完成，请使用标准CUDA同步机制（如事件）或等待流。

6. C++ API 的内存管理

TensorRT提供了两种机制来允许应用程序对设备内存进行更多的控制。

默认情况下，在创建IExecutionContext时，会分配持久设备内存来保存激活数据。为避免这个分配，请调用createExecutionContextWithoutDeviceMemory。然后应用程序会调用IExecutionContext :: setDeviceMemory()来提供运行网络所需的内存。内存块的大小由ICudaEngine :: getDeviceMemorySize()返回。

此外，应用程序可以通过实现IGpuAllocator接口来提供在构建和运行时使用的自定义分配器。实现接口后，请调用：

setGpuAllocator(&allocator);

在IBuilder或IRuntime接口上。所有的设备内存都将通过这个接口来分配和释放。

7. 调整engine

TensorRT可以为一个engine装填新的weights，而不用重新build。

1. 在build之前申请一个可refittable的engine

...
builder->setRefittable(true); 
builder->buildCudaEngine(network);

2. 创建一个refitter对象

ICudaEngine* engine = ...;
IRefitter* refitter = createInferRefitter(*engine,gLogger)

3. 更新你想更新的weights
如：为一个叫MyLayer的卷积层更新kernel weights

Weights newWeights = ...;
refitter.setWeights("MyLayer",WeightsRole::kKERNEL,
                    newWeights);

这个新的weights要和原始用来build engine的weights具有相同的数量。
setWeights出错时会返回false。

4. 找出哪些weights需要提供。
这通常需要调用两次IRefitter::getMissing，第一次调用得到Weights对象的数目，第二次得到他们的layers和roles。

const int n = refitter->getMissing(0, nullptr, nullptr);
std::vector<const char*> layerNames(n);
std::vector<WeightsRole> weightsRoles(n);
refitter->getMissing(n, layerNames.data(), 
                        weightsRoles.data());

5. 提供missing的weights（顺序无所谓）

for (int i = 0; i < n; ++i)
    refitter->setWeights(layerNames[i], weightsRoles[i],
                         Weights{...});

只需提供missing的weights即可，如果提供了额外的weights可能会触发更多weights的需要。

6. 更新engine

bool success = refitter->refitCudaEngine();
assert(success);

如果success的值为false，可以检查一下diagnostic log，也许有些weights还是missing的。

7. 销毁refitter

refitter->destroy();

如果想要查看engine中所有可重新调整的权重，可以使用refitter-> getAll(...)，类似于步骤4中的如何使用getMissing。

参考

DEEP LEARNING SDK DOCUMENTATION