4. TensorRT引擎构建基础

TensorRT 的推理部署并不是“加载模型直接跑”这么简单。它本质上是一个离线编译 + 在线执行的两阶段体系：先把训练框架导出的模型（ONNX、Caffe、UFF 等）翻译成 TensorRT 自己的计算图，再针对目标 GPU 架构做层融合、精度校准、内核自动调优，最终生成一个序列化的 Engine 文件；运行时再反序列化出 Engine，创建 Context 执行推理。理解这条链路上的核心对象，是写出可维护、可复现部署代码的前提。

一、六个核心对象的职责边界

TensorRT 的 C++/Python API 中，构建阶段和执行阶段各有一组对象，它们的关系可以用一句话概括：Builder 用 Network 描述计算图，用 Config 规定优化策略，产出 Engine；Engine 在运行时派生出 Context 承载实际推理状态。

  • Builder（IBuilder）：构建阶段的“总调度”。它负责创建 Network、创建 BuilderConfig、调用 buildSerializedNetwork() 或 buildEngineWithConfig() 触发优化与序列化。Builder 本身不保存模型信息，只提供工厂能力。
  • Network（INetworkDefinition）：模型的计算图表示。它由 Tensor（张量）和 Layer（层）组成，记录每个层的类型、输入输出张量、维度、数据类型。Parser 解析 ONNX 后，结果就写进 Network。
  • Parser（如 ONNXParser）：模型格式的“翻译器”。它读取 ONNX 文件，把 protobuf 描述的算子图映射为 Network 中的 Layer 与 Tensor，并报告不支持的算子。
  • Config（IBuilderConfig）：构建的“策略面板”。它决定精度模式（FP32/FP16/INT8）、workspace 显存上限、是否启用 TF32、INT8 校准器、profiling 粒度、硬件兼容性等。同一个 Network 配不同 Config，会得到性能差异明显的 Engine。
  • Engine（ICudaEngine）：构建的最终产物，是针对特定 GPU 架构和 TensorRT 版本优化后的可执行计划。它包含融合后的层、选定的内核、张量内存布局，可序列化为 .plan 文件持久化。
  • Context（IExecutionContext）：Engine 的“运行实例”。它持有推理所需的中间激活内存、绑定（binding）地址、流（stream）等状态。一个 Engine 可以创建多个 Context 并发执行，但每个 Context 需要独立的显存缓冲。
关键区分：Engine 是只读的、可共享的优化产物；Context 是有状态的、不可共享的执行上下文。多线程推理时，应共享 Engine、每线程一个 Context。

二、构建流程的对象流转

下面这张图展示了从 ONNX 文件到推理输出的完整对象流转关系，注意区分“构建期”和“运行期”两个阶段，以及序列化文件在其中的桥梁作用。

构建期（离线，可跨机器） 运行期（在线，绑定目标 GPU） ONNX 文件 model.onnx Parser ONNXParser Network INetworkDefinition Builder IBuilder Config IBuilderConfig createConfig Engine 文件 model.plan buildSerializedNetwork Engine ICudaEngine Context IExecutionContext CUDA Stream 异步执行 反序列化 createExecutionContext enqueueV3 输入张量 input: [1,3,224,224] 输出张量 output: [1,1000]

三、最小可运行构建流程

下面给出一段 Python 最小示例，目标是把一个 ONNX 分类模型构建为 FP16 Engine 并序列化。示例参数为教学假设值，实际使用时需按模型和 GPU 调整。

import tensorrt as trt

logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)

# 1. 创建 Network，显式指定 batch 维度
network = builder.create_network(
    1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

# 2. 创建 Parser 并解析 ONNX
parser = trt.OnnxParser(network, logger)
with open("resnet18.onnx", "rb") as f:
    if not parser.parse(f.read()):
        for i in range(parser.num_errors):
            print(parser.get_error(i))
        raise RuntimeError("ONNX 解析失败")

# 3. 创建 Config，设置精度与 workspace
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1 GiB
if builder.platform_has_fast_fp16:
    config.set_flag(trt.BuilderFlag.FP16)

# 4. 构建并序列化 Engine
serialized = builder.build_serialized_network(network, config)
with open("resnet18.plan", "wb") as f:
    f.write(serialized)
print("Engine 构建完成，大小:", len(serialized) / 1024 / 1024, "MB")

运行这段代码后，如果终端没有报错并打印出 Engine 大小，说明构建成功。此时 resnet18.plan 已经是一个与当前 GPU 架构绑定的优化产物。接下来在运行期加载它：

runtime = trt.Runtime(logger)
with open("resnet18.plan", "rb") as f:
    engine = runtime.deserialize_cuda_engine(f.read())
context = engine.create_execution_context()

# 查看绑定张量
for i in range(engine.num_io_tensors):
    name = engine.get_tensor_name(i)
    mode = engine.get_tensor_mode(name)
    shape = engine.get_tensor_shape(name)
    dtype = engine.get_tensor_dtype(name)
    print(name, mode, shape, dtype)

四、构建结果如何判断

构建是否“成功”不能只看有没有抛异常，还要关注几个可观察信号：

  • Parser 错误数：parser.num_errors 必须为 0。若 ONNX 含 TensorRT 不支持的算子（如某些自定义算子），解析会失败，需要先做算子替换或插件实现。
  • Engine 大小与层数：构建后可用 engine.num_layers 查看融合后的层数。通常远小于原始 ONNX 节点数，说明融合生效。
  • 精度标志是否生效：若设置了 FP16 但 builder.platform_has_fast_fp16 为 False，则 FP16 不会真正启用，需检查 GPU 计算能力。
  • workspace 是否够用：workspace 不足时构建可能失败或回退到较慢的内核。可先给较大值（如 2 GiB）构建成功，再逐步下调观察性能变化。
  • 序列化文件可反序列化：用 runtime.deserialize_cuda_engine() 能成功加载，且 engine.num_io_tensors 与预期一致，才算闭环。
Engine 与 GPU 架构、TensorRT 版本强绑定。在 A100 上构建的 .plan 不能直接在 T4 上反序列化使用，必须重新构建。部署时建议在目标机型上构建，或使用 BuilderConfig 的硬件兼容性选项（会牺牲部分性能）。

五、对象生命周期与常见误区

构建完成后，builder、network、parser、config 都可以释放，它们不参与推理。运行期真正需要长期持有的是 engine 和 context。常见误区有三个：

  1. 把 Engine 当模型文件反复构建：每次启动都重新 build 会浪费几十秒到几分钟。正确做法是首次构建后序列化，后续直接反序列化。
  2. 多线程共享一个 Context：Context 内部有推理状态，多线程同时调用会导致结果错乱或崩溃。应每线程创建独立 Context，并各自分配输入输出显存。
  3. 忽略 EXPLICIT_BATCH 标志：不设置该标志时，Network 使用隐式 batch 语义，动态 shape 和部分算子行为会受限。现代 TensorRT 建议始终显式启用。
← 上一章 📖 返回目录 下一章 →