3、TensorRT核心API入门：创建Builder与Network、定义输入输出张量、添加层（Layer）、构建Engine与序列化。

好，咱们直接进入正题。这一章是TensorRT插件开发的基础中的基础。说白了，你后面写的所有自定义算子，最终都要嵌入到这套流程里。我刚开始接触TensorRT时，也踩过不少坑，尤其是Builder和Network的配置，稍不注意就性能翻车。今天咱们就把这几个核心API捋一遍。

3.1 创建Builder与Network

Builder，顾名思义，就是负责构建Engine的“工头”。它决定了你最终生成的推理引擎长什么样。Network呢，就是你的计算图——你定义的所有层、张量、连接关系，都塞进这个Network里。

创建Builder很简单，就一行代码：

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

这里有个细节：gLogger 是日志记录器。我建议你至少用 ILogger::Severity::kWARNING 级别，不然调试时信息太少，出了问题都不知道错在哪。

有了Builder，接下来创建Network：

const auto explicitBatch = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
nvinfer1::INetworkDefinition* network = builder->createNetworkV2(explicitBatch);

嗯，这里要注意。为什么用 createNetworkV2 而不是老版本的 createNetwork？因为V2支持显式batch模式。你想想看，如果你的模型输入batch size是动态的，或者你想用TensorRT 8.x以上的新特性，就必须用V2。我曾经在迁移旧项目时，忘了改这个，结果跑出来的结果全是错的，排查了半天。

3.2 定义输入输出张量

Network创建好了，接下来要告诉它：你的模型吃进去什么，吐出什么。这就是定义输入输出张量。

定义输入张量：

nvinfer1::ITensor* input = network->addInput("input", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{1, 3, 224, 224});

这里 Dims4 表示四维张量，分别是NCHW格式。我个人习惯把维度写清楚，哪怕有些维度是1，也不省略。为什么？因为后续添加层时，维度推导容易出错，写清楚能省很多调试时间。

定义输出张量稍微特殊一点——输出张量通常不是直接addInput，而是从某个层的输出中指定：

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

注意，markOutput 之前，你最好确认一下这个张量的数据类型和维度。我遇到过一个问题：某个层的输出是FP16，但我忘了设置，结果Builder自动转成了FP32，性能直接掉了一半。

3.3 添加层（Layer）

这是最核心的部分。TensorRT提供了几十种预定义层，比如卷积、池化、全连接、激活函数等。添加层的方式很统一：

nvinfer1::IConvolutionLayer* conv = network->addConvolutionNd(*input, 64, nvinfer1::DimsHW{3, 3}, kernelWeights, biasWeights);
conv->setStrideNd(nvinfer1::DimsHW{1, 1});
conv->setPaddingNd(nvinfer1::DimsHW{1, 1});

这里 kernelWeights 和 biasWeights 是权重数据，需要提前用 Weights 结构体包装好。我刚开始写时，总忘记把权重数据从模型文件里读出来，结果Builder直接报错。

添加层的顺序很重要。TensorRT会按照你添加的顺序构建计算图。但注意，它内部会做图优化，比如融合相邻的Conv+ReLU。所以你不必担心顺序影响性能，但必须保证数据流的逻辑正确。

举个例子，如果你想加一个ReLU激活层：

nvinfer1::IActivationLayer* relu = network->addActivation(*conv->getOutput(0), nvinfer1::ActivationType::kRELU);

你看，上一层的输出直接作为下一层的输入。这种链式调用很直观，但容易犯一个错误：忘记保存中间层的指针。我建议你把每个关键层的指针都存下来，方便后续调试或修改参数。

3.4 构建Engine与序列化

Network搭好了，接下来就是“编译”成Engine。这一步是TensorRT最核心的优化过程。它会做算子融合、内存复用、内核调优等。

构建Engine需要先创建BuilderConfig：

nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 1 << 30); // 1GB
config->setFlag(nvinfer1::BuilderFlag::kFP16); // 启用FP16推理

这里 setMemoryPoolLimit 设置工作空间大小。我建议根据你的GPU显存来设，别设太大，否则多个模型同时跑时会OOM。也别设太小，否则优化效果差。

然后构建Engine：

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

注意，我用了 buildSerializedNetwork 而不是 buildEngineWithConfig。为什么？因为前者直接返回序列化后的模型，你可以保存到文件，下次直接加载，省去重新构建的时间。你想想看，一个复杂的模型构建可能要几分钟，而反序列化只需要几秒钟。

序列化后的数据可以保存到文件：

std::ofstream file("model.engine", std::ios::binary);
file.write(reinterpret_cast<const char*>(serializedModel->data()), serializedModel->size());
file.close();

反序列化加载时：

nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(gLogger);
nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(serializedModel->data(), serializedModel->size());

3.5 避坑指南与个人经验

最后，分享几个我实际项目中踩过的坑：

• BuilderConfig的精度设置： 如果你启用了FP16，但模型中有某些层对精度敏感（比如BatchNorm），记得用 config->setFlag(nvinfer1::BuilderFlag::kPREFER_PRECISION_CONSTRAINTS) 来强制某些层使用FP32。
• 动态shape： 如果你的输入shape是动态的，需要在BuilderConfig中设置优化profile。我建议至少设置三个profile：最小、最优、最大。否则推理时输入尺寸一变，Engine就罢工。
• 内存泄漏： TensorRT的API对象（Builder、Network、Config等）都需要手动释放。我习惯用智能指针或者RAII机制来管理，省得忘了delete。

好了，这一章的内容就到这里。核心API的用法其实不难，难的是理解每个参数背后的含义。下一章我们会深入自定义算子的编写，到时候这些API会反复用到。建议你把今天的代码亲手敲一遍，跑通一个最简单的模型（比如一个Conv+ReLU），感受一下整个流程。