3. TensorRT核心概念与工作流程

TensorRT 并不是一个“直接运行 ONNX 的推理框架”，而是一套把训练好的模型编译成针对特定 GPU 高度优化的推理引擎的工具链。理解它的关键，是先把几个核心概念串成一条线：计算图、层与张量、精度模式、构建器与运行时、序列化与反序列化，以及 ONNX 解析与模型导入。只有把这条线走通，后面在 Jetson 上做 FP16/INT8 加速、显存优化和吞吐调优时，才知道每一步在改什么、为什么有效。

计算图、层与张量：TensorRT 眼中的模型

一个神经网络在 TensorRT 内部被表示为一张有向无环计算图。图中的节点是层（Layer），边是张量（Tensor）。层负责运算，例如卷积、矩阵乘、激活、池化、归一化；张量负责承载数据，它带有名称、维度（shape）、数据类型（dtype）和内存布局。以常见的 ResNet 类分类网络为例，输入张量形如 [N, 3, 224, 224]，经过若干卷积层后张量维度逐步缩小、通道数增大，最后接一个全连接或全局池化层输出 [N, 1000] 的 logits。TensorRT 在构建阶段会遍历这张图，做层融合、常量折叠、精度选择和张量内存复用，最终生成一个引擎。

这里要区分两个概念：网络定义（INetworkDefinition）是构建期的可变图，可以添加、修改层；引擎（ICudaEngine）是构建后的固化产物，结构、精度和内存规划都已确定，运行时只负责按绑定（binding）喂数据、取结果。张量在引擎中通常分为输入、输出和内部张量，内部张量由 TensorRT 自行分配和复用，用户一般只关心输入输出绑定。

TensorRT 构建与运行流程：从 ONNX 到推理引擎 ONNX 模型 .onnx 文件 含计算图与权重 ONNX Parser 解析图结构 映射为 TRT 层 网络定义 INetworkDefinition 层 + 张量（可变） 构建器 Builder 层融合 / 精度选择 内存规划 / 内核调优 推理引擎 ICudaEngine 固化结构 + 精度 序列化 engine.plan 文件 反序列化 加载到显存 运行时 Runtime IExecutionContext 推理执行 enqueueV3 / execute 输入 / 输出绑定 显存指针 + shape 精度模式对比（构建期选择，影响速度、显存与精度） FP32 32 位浮点，精度最高 速度基准，显存占用最大 适合精度敏感任务 FP16 16 位浮点，速度约 1.5~2× 显存减半，精度损失小 Jetson 上首选加速档 INT8 8 位整型，速度约 2~4× 需校准集确定缩放因子 精度需验证，适合高吞吐

精度模式：FP32、FP16 与 INT8 的取舍

精度模式是 TensorRT 优化中最直接的“加速旋钮”。FP32 是默认精度，数值稳定、精度最高，但计算和显存开销最大。FP16 把权重和激活降到 16 位浮点，在 Jetson 的 GPU 上通常能获得约 1.5～2 倍的速度提升，显存占用减半，而多数视觉模型的精度损失很小，因此是 Jetson 部署的首选档位。INT8 进一步把数据压到 8 位整型，速度可达 FP32 的 2～4 倍，但它不是简单地把浮点截断，而是需要校准（calibration）：用一批有代表性的输入数据统计每层激活的数值分布，确定缩放因子（scale）和零点（zero point），从而在量化后尽量保持输出一致。

在代码层面，精度通过构建配置指定。例如启用 FP16 只需设置一个标志位：

config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.FP16)   # 启用 FP16
# config.set_flag(trt.BuilderFlag.INT8) # 启用 INT8，需配合校准器
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB workspace
engine = builder.build_engine(network, config)

需要注意的是，精度标志是“允许”而非“强制”。TensorRT 会逐层判断该层是否支持目标精度，若某层不支持则自动回退到更高精度，因此最终引擎可能是混合精度。判断是否真正生效，可以查看构建日志中类似 Conv_0 ... FP16 的层信息，或对比不同精度下引擎的推理耗时与输出误差。

构建器与运行时：两个阶段、两种角色

TensorRT 的工作流明确分为构建期和运行期。构建期由 IBuilder 主导，负责解析模型、优化图、选择精度、分配工作空间并生成引擎。这个过程计算密集、耗时较长，在 Jetson 上构建一个中等规模模型可能需要数十秒到数分钟，因此不应在每次启动时都重新构建。运行期由 IRuntime 和 IExecutionContext 主导，负责加载已序列化的引擎、绑定输入输出显存、执行推理。运行时不做图优化，启动快、开销小，适合部署场景。

一个常见的工程做法是：在开发机或 Jetson 上一次性构建引擎并序列化保存，部署时只做反序列化。序列化把引擎保存为 .plan 文件，反序列化则从该文件恢复引擎。需要特别注意的是，引擎与构建时的 GPU 架构、TensorRT 版本、精度配置强绑定，换设备或升级 TensorRT 后通常需要重新构建，否则反序列化可能失败或行为异常。

ONNX 解析与模型导入

TensorRT 原生不直接读取 PyTorch 或 TensorFlow 模型，主流路径是先把模型导出为 ONNX，再用 ONNX Parser 导入。导入时 Parser 会遍历 ONNX 图的节点，把支持的算子映射为 TensorRT 层，不支持的算子则报错或需要自定义插件（Plugin）。因此模型导入阶段最常见的两类问题是：算子不支持和动态维度处理不当。

以 PyTorch 导出为例，典型流程如下：

# 1. 导出 ONNX（在 PyTorch 侧执行）
torch.onnx.export(model, dummy_input, "model.onnx",
                  input_names=["input"], output_names=["output"],
                  opset_version=13, dynamic_axes={"input": {0: "batch"}})

# 2. TensorRT 导入并构建
import tensorrt as trt
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)
with open("model.onnx", "rb") as f:
    if not parser.parse(f.read()):
        for i in range(parser.num_errors):
            print(parser.get_error(i))

如果解析失败，parser.get_error(i) 会给出具体算子或维度问题，这是排查导入错误的第一手信息。对于动态 batch，需要在构建配置中设置优化 profile，指定最小、最优、最大 shape，TensorRT 会针对该范围生成可复用的引擎。若模型包含 Parser 不支持的算子，则需要编写自定义插件并在网络中手动添加该层，这是进阶部署中绕不开的一环。

在 Jetson 上构建引擎时，建议先用小 batch（如 1）和 FP16 快速验证流程是否跑通，再逐步调整 batch、精度和 workspace 大小。构建日志中的层信息、显存占用和耗时，是判断优化是否生效的直接依据。

把上述概念串起来看：ONNX 提供模型描述，Parser 把它翻译成 TensorRT 的网络定义，Builder 在指定精度下优化并生成引擎，序列化把引擎落盘，Runtime 在部署时反序列化并执行推理。理解这条链路中每一步的输入、输出和约束，才能在 Jetson 上做出既快又稳的推理部署。

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