4、TensorRT API入门：创建Builder与Network、配置Builder参数（工作空间、精度）、构建Engine与序列化

好，咱们正式开始动手写代码了。

这一章，我带你走一遍 TensorRT 推理最核心的流程：创建 Builder → 定义 Network → 配置参数 → 构建 Engine → 序列化保存。说白了，就是教会 TensorRT 怎么“看懂”你的模型，然后把它变成一张能飞快运行的“蓝图”。

我个人习惯把这一步叫做“炼丹前的炉灶搭建”。炉灶搭不好，后面火候再大也白搭。

4.1 创建 Builder 与 Network

先上最基础的代码骨架。你想想看，TensorRT 本身是个 C++ 库，但咱们做工程时 Python 用得更多。所以我这里以 Python API 为例，C++ 的接口逻辑完全一样，只是语法不同。

import tensorrt as trt

# 创建 Builder
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)

# 创建 Network
network = builder.create_network(
    1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
)

# 创建配置对象
config = builder.create_builder_config()

嗯，这里要注意几个点。

第一，Logger 的级别。 我刚开始用的时候，图省事设成了 trt.Logger.INFO，结果控制台刷屏一样输出各种信息，找关键错误像大海捞针。后来我学乖了，开发调试用 INFO，正式部署一律用 WARNING 或 ERROR。你想想看，生产环境谁有空看那些“加载了某个算子”的日志？

第二，EXPLICIT_BATCH 这个 flag。 这是 TensorRT 8.x 之后强烈推荐的模式。以前的老模式（隐式 Batch）会把 batch size 硬编码到网络中，换 batch size 就得重新构建 Engine，非常死板。显式 Batch 模式下，batch 维度被当作一个普通维度处理，灵活得多。我曾经在迁移一个老项目时，忘了加这个 flag，结果跑推理时 batch size 一变化就报错，排查了半天……嗯，从那以后我每次创建 Network 都会默念一遍：显式、显式、显式。

4.2 配置 Builder 参数

Builder 建好了，Network 也搭起来了，接下来就是告诉 TensorRT：“嘿，你想怎么优化这个模型？”

配置参数主要就两样：工作空间（Workspace） 和 精度（Precision）。

4.2.1 工作空间（Workspace）

工作空间，说白了就是 TensorRT 在构建 Engine 时能用的“临时内存上限”。它决定了 TensorRT 能尝试多少种算子融合策略。空间越大，它越有可能找到更优的融合方案，推理速度可能更快。

# 设置工作空间大小，单位是字节
# 这里设置为 1GB
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)

我个人习惯给个 1GB 到 2GB。但要注意，这不是说你显卡显存有 8GB 就能设 8GB。构建时的内存消耗和推理时的显存占用是两码事。我曾经在一个边缘设备上（只有 4GB 内存），把工作空间设成了 3GB，结果构建 Engine 时直接 OOM 了。后来我改成 512MB，虽然构建时间长了点，但至少能跑起来。

4.2.2 精度设置

精度是 TensorRT 加速的核心手段之一。默认是 FP32，但你可以开启 FP16 或 INT8。

# 开启 FP16
if builder.platform_has_fast_fp16:
    config.set_flag(trt.BuilderFlag.FP16)
    print("FP16 已开启，速度起飞！")
else:
    print("你的显卡不支持 FP16，老老实实用 FP32 吧。")

# 开启 INT8（需要校准数据集）
# if builder.platform_has_fast_int8:
#     config.set_flag(trt.BuilderFlag.INT8)
#     # 还需要设置校准器，后面章节会细讲

这里有个坑：platform_has_fast_fp16 这个检查一定要做。我见过有人不管三七二十一直接设 FP16，结果在旧款 Tesla 显卡上跑，速度反而比 FP32 还慢。为什么？因为那些显卡的 FP16 计算单元是模拟的，没有硬件加速。

4.3 构建 Engine 与序列化

参数配好了，Network 也定义完了，接下来就是见证奇迹的时刻——构建 Engine。

4.3.1 构建 Engine

# 序列化模型（构建 Engine）
serialized_engine = builder.build_serialized_network(network, config)

# 或者你也可以先构建 Engine，再序列化
# engine = builder.build_engine(network, config)
# serialized_engine = engine.serialize()

我个人更推荐第一种方式：build_serialized_network。一步到位，直接拿到序列化后的字节流。省去了先构建 Engine 再序列化的中间步骤，代码更简洁。

构建 Engine 这一步，是 TensorRT 最耗时的地方。一个复杂的模型（比如 ResNet-152 或 BERT-Large），构建时间可能长达几分钟甚至十几分钟。你想想看，TensorRT 在这段时间里在干嘛？它在疯狂地尝试各种算子融合、内存复用、层合并……说白了，就是在“编译”你的模型。

4.3.2 序列化与反序列化

构建好的 Engine 可以直接用，但更常见的做法是把它序列化保存到硬盘上。这样下次启动程序时，直接加载序列化文件，省去了重新构建的时间。

# 保存到文件
with open("model.trt", "wb") as f:
    f.write(serialized_engine)

# 加载时反序列化
with open("model.trt", "rb") as f:
    serialized_engine = f.read()

runtime = trt.Runtime(logger)
engine = runtime.deserialize_cuda_engine(serialized_engine)

这里有个小细节：序列化文件是平台相关的。你在 RTX 3090 上构建的 Engine，不能直接拿到 Jetson Nano 上用。因为不同 GPU 的架构不同（比如 Ampere 和 Turing），TensorRT 在构建时会针对特定架构做优化。我曾经犯过这个错，在开发机上构建好 Engine，拷到客户现场的 Tesla T4 上加载，结果直接报错……嗯，从那以后我每次部署都会在目标机器上重新构建一次。

好了，这一章的内容就到这里。你跟着代码走一遍，应该就能跑通一个最简单的 TensorRT 推理流程了。下一章，咱们聊聊怎么把 PyTorch 模型转成 TensorRT 能吃的格式——ONNX。那又是一个充满坑与经验的故事。