4、TensorRT基础API：创建Builder与Config、构建Network与Engine、序列化与反序列化

好，咱们进入第四章。这一章是TensorRT的“地基”。

说白了，你后面所有花里胡哨的优化技巧，都得靠这几个API搭起来。我刚开始接触TensorRT时，觉得这玩意儿不就是个推理引擎嘛，直接扔个模型进去不就完了？结果一上手，发现Builder、Config、Network、Engine这几个概念绕得我头晕。

嗯，今天咱们就把它们彻底捋清楚。

4.1 创建Builder与Config：一切从这里开始

Builder，翻译过来就是“构建者”。它是TensorRT的入口点。你想想看，你要盖一栋楼，总得先有个施工队吧？Builder就是这个施工队的队长。

创建Builder的代码很简单：

#include <NvInfer.h>
using namespace nvinfer1;

// 创建Builder
IBuilder* builder = createInferBuilder(gLogger);

这里有个细节——gLogger。它是一个日志记录器，负责输出TensorRT的日志信息。我个人习惯用Logger类继承ILogger，然后重写log()方法。这样你能控制哪些日志需要显示，哪些可以忽略。

小技巧： 生产环境中，我建议把WARNING级别以上的日志都记录下来。调试时再打开INFO级别。否则日志刷屏，你根本找不到关键信息。

有了Builder，下一步就是创建Config。Config是Builder的配置对象，它决定了Builder怎么干活。

// 创建Config
IBuilderConfig* config = builder->createBuilderConfig();

Config能配置的东西很多，比如：

  • 工作空间大小：setMemoryPoolLimit()，控制显存使用上限
  • 精度模式：FP32、FP16、INT8
  • 动态形状：允许输入尺寸变化
  • DLA加速：如果你有NVIDIA的DLA硬件

举个例子，设置FP16模式：

if (builder->platformHasFastFp16()) {
    config->setFlag(BuilderFlag::kFP16);
    std::cout << "FP16模式已开启" << std::endl;
}

这里要注意：platformHasFastFp16()会检查你的GPU是否支持FP16加速。我曾经在T4上跑得好好的，换到K80上直接报错——因为K80不支持FP16。所以，别偷懒，一定要加这个检查。

避坑指南： 我曾经在生产环境中忘记设置工作空间大小，结果模型在推理时频繁OOM。后来我养成了习惯：每次创建Config后，第一件事就是设置setMemoryPoolLimit()。建议设置为显存总量的70%-80%，留点余量给其他进程。

4.2 构建Network与Engine：模型的核心

Network，就是你的神经网络。它由一系列层（Layer）和张量（Tensor）组成。Engine，则是优化后的可执行对象。你可以把Network理解成“设计图”，Engine理解成“盖好的房子”。

创建Network的代码：

// 创建Network
INetworkDefinition* network = builder->createNetworkV2(0U);

参数0U表示使用默认的网络标志。如果你想支持动态形状，需要传1U << static_cast<int>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH)。这个后面会细讲。

有了Network，你就可以往里面添加层了。比如添加一个卷积层：

// 添加输入张量
ITensor* input = network->addInput("input", DataType::kFLOAT, Dims4{1, 3, 224, 224});

// 添加卷积层
auto conv = network->addConvolutionNd(*input, 64, DimsHW{3, 3}, weightMap["conv.weight"], weightMap["conv.bias"]);
conv->setStrideNd(DimsHW{1, 1});
conv->setPaddingNd(DimsHW{1, 1});

// 标记输出
network->markOutput(*conv->getOutput(0));

嗯，这里要注意：TensorRT的API是C++风格的，但很多操作符（比如ReLU、池化）都有对应的addXxx()方法。你不需要手动实现前向传播，TensorRT会帮你搞定。

Network构建完成后，就该生成Engine了：

// 构建Engine
IHostMemory* serializedModel = builder->buildSerializedNetwork(*network, *config);
IRuntime* runtime = createInferRuntime(gLogger);
ICudaEngine* engine = runtime->deserializeCudaEngine(serializedModel->data(), serializedModel->size());

这里我用了buildSerializedNetwork()而不是buildEngineWithConfig()。为什么？因为前者返回序列化后的模型，方便保存到磁盘。后者直接返回Engine，但无法持久化。

核心要点： Engine是TensorRT优化的最终产物。它包含了所有层的融合、内存优化、内核选择等结果。一旦生成，就可以反复使用，不需要重新构建。

4.3 序列化与反序列化：让模型“活”起来

序列化，就是把Engine保存到磁盘。反序列化，就是从磁盘加载Engine。这俩操作，说白了就是“存”和“取”。

序列化的代码：

// 序列化Engine到文件
std::ofstream file("model.engine", std::ios::binary);
file.write(static_cast<const char*>(serializedModel->data()), serializedModel->size());
file.close();

反序列化的代码：

// 从文件反序列化Engine
std::ifstream file("model.engine", std::ios::binary);
file.seekg(0, std::ios::end);
size_t size = file.tellg();
file.seekg(0, std::ios::beg);

std::vector<char> buffer(size);
file.read(buffer.data(), size);

IRuntime* runtime = createInferRuntime(gLogger);
ICudaEngine* engine = runtime->deserializeCudaEngine(buffer.data(), size);

这里有个坑：序列化后的Engine是平台相关的。什么意思？你在T4上生成的Engine，拿到A100上可能跑不了。因为TensorRT会针对特定GPU架构生成最优的CUDA内核。

避坑指南： 我曾经在开发机上用RTX 3090生成Engine，然后部署到服务器上的T4上。结果反序列化直接报错。后来我学乖了：在目标设备上生成Engine，或者使用kSAFE标志生成跨平台Engine（但性能会打折扣）。

另外，序列化后的Engine文件大小通常比原始模型小很多。比如一个ResNet-50的ONNX模型大约100MB，序列化后可能只有30MB。这是因为TensorRT做了权重压缩和层融合。

4.4 完整流程：从Builder到Engine

咱们把整个流程串起来，看看一个完整的构建过程：

// 1. 创建Builder
IBuilder* builder = createInferBuilder(gLogger);

// 2. 创建Config
IBuilderConfig* config = builder->createBuilderConfig();
config->setMemoryPoolLimit(MemoryPoolType::kWORKSPACE, 1 << 30); // 1GB

// 3. 创建Network
INetworkDefinition* network = builder->createNetworkV2(0U);

// 4. 添加网络结构（这里省略具体层）
// ...

// 5. 构建序列化模型
IHostMemory* serializedModel = builder->buildSerializedNetwork(*network, *config);

// 6. 创建Runtime并反序列化
IRuntime* runtime = createInferRuntime(gLogger);
ICudaEngine* engine = runtime->deserializeCudaEngine(serializedModel->data(), serializedModel->size());

// 7. 清理资源
serializedModel->destroy();
network->destroy();
config->destroy();
builder->destroy();

嗯，这里要注意资源管理。TensorRT的API是C风格的，没有RAII。所以你得手动调用destroy()。我建议用智能指针封装一下，或者写个RAII包装类。否则内存泄漏会让你头疼。

个人经验： 我习惯把Builder、Config、Network、Runtime、Engine都封装成std::unique_ptr，自定义删除器。这样即使异常抛出，资源也能自动释放。代码看起来也更干净。

4.5 总结

这一章咱们把TensorRT的四个核心API捋了一遍：

  • Builder：施工队长，负责统筹全局
  • Config：施工方案，控制精度、内存等
  • Network：设计图，定义网络结构
  • Engine：盖好的房子，可执行推理

序列化和反序列化，就是让Engine能存能取，方便部署。

下一章，咱们会深入Network的构建细节，看看如何从ONNX模型导入，以及如何手动搭建网络。到时候，你会发现自己已经能掌控TensorRT的“命脉”了。

← 上一章 📖 返回目录 下一章 →