📚 课程目录第3章 TensorRT推理引擎基础

3. TensorRT推理引擎基础

在Jetson平台上，TensorRT并不是一个“把模型跑起来”的黑盒，而是一套围绕推理图做编译优化的运行时系统。理解它的架构与工作流，是后续做INT8校准、动态形状调优和显存规划的前提。本节围绕四个核心问题展开：TensorRT由哪些部分组成、模型如何被解析与序列化、ONNX如何转换成TensorRT引擎、动态形状与精度校准在转换阶段如何介入。

一、TensorRT的架构分层

TensorRT的运行时结构可以自上而下分为四层。最上层是模型导入层，负责接收ONNX、Caffe、UFF等格式，其中ONNX是当前Jetson部署的主流入口。第二层是图优化与构建层，由Builder和Network Definition组成，负责层融合、精度选择、内核自动调优和显存规划。第三层是序列化层，把构建好的引擎固化为plan文件。最底层是执行层，由Runtime、ExecutionContext和CUDA内核组成，负责在设备上实际调度推理。

需要特别区分两个概念：Builder阶段在主机或设备上完成，耗时较长，产出的是不可变的引擎；Runtime阶段只负责加载引擎并执行，不参与图优化。因此部署时通常把构建放在开发机或首次启动时完成，把plan文件随应用分发，避免每次启动都重新编译。

二、模型解析与序列化

ONNX Parser是TensorRT读取模型的第一道关口。它把ONNX的计算图逐节点翻译成TensorRT的INetworkDefinition，这个过程称为解析。解析时Parser会检查算子是否被支持、属性是否完整、输入输出维度是否合法。如果遇到不支持的算子，构建会直接失败并给出节点名称，此时需要自定义插件或修改模型结构。

解析完成后进入构建阶段。Builder会做几件关键的事：把Conv+BN+ReLU这类连续算子融合成一个内核，减少显存访问；根据目标精度选择FP16或INT8的Tensor Core路径；对每个层尝试多种内核实现并实测耗时，选出最快的方案。这个过程在Jetson上可能持续几十秒到几分钟，取决于模型规模。

构建成功后得到的是序列化引擎。通过builder.build_serialized_network()或engine.serialize()可以把引擎写成plan文件。plan文件与具体的TensorRT版本、GPU架构、精度配置强绑定，换设备或升级版本后通常需要重新构建。反序列化时用runtime.deserialize_cuda_engine()加载，再创建ExecutionContext即可执行推理。

# 构建并序列化引擎（Python API 示意）
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:
    parser.parse(f.read())

config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)
config.set_flag(trt.BuilderFlag.FP16)

# 动态形状：设置优化 profile
profile = builder.create_optimization_profile()
profile.set_shape("input",
                  min=(1, 3, 224, 224),
                  opt=(4, 3, 416, 416),
                  max=(8, 3, 640, 640))
config.add_optimization_profile(profile)

serialized = builder.build_serialized_network(network, config)
with open("model.plan", "wb") as f:
    f.write(serialized)

三、ONNX到TensorRT的转换流程

完整的转换流程可以概括为五步：导出ONNX、检查算子兼容性、配置Builder参数、构建引擎、验证输出一致性。其中最容易出问题的是前两步。

导出ONNX时，PyTorch的torch.onnx.export需要指定opset_version，建议使用11或更高版本以获得更好的算子覆盖。如果模型包含动态维度，要在dynamic_axes中显式声明，例如把batch维和空间维标记为动态。导出后可以用onnx.checker.check_model()做基本校验，再用polygraphy surgeon sanitize或onnxsim做常量折叠和冗余节点清理，减少Parser的负担。

算子兼容性检查可以借助trtexec --onnx=model.onnx --verbose，它会列出每个节点的解析结果。如果出现Unsupported operator，需要确认该算子在当前TensorRT版本中是否已支持，或者用插件替代。Jetson上常见的坑包括自定义的Resize、GridSample、NonMaxSuppression等，这些在较新版本中已有原生支持，但旧版本仍需插件。

构建完成后必须做数值一致性验证。用同一组输入分别跑ONNX Runtime和TensorRT，比较输出的最大绝对误差和相对误差。FP16下通常允许1e-3量级的差异，INT8下会更大，需要结合业务指标判断是否可接受。如果误差超标，优先检查校准集是否具有代表性，再考虑回退到FP16。

四、动态形状与精度校准的介入点

动态形状通过Optimization Profile控制。每个profile定义一组输入张量的min/opt/max形状，Builder会针对opt形状做内核调优，同时保证min到max范围内都能正确执行。如果实际输入超出max，运行时会报错，必须重建引擎。因此设置max时要留足余量，但也不能过大，否则会浪费显存并降低调优效果。

INT8校准则在构建阶段通过IInt8EntropyCalibrator2接口介入。校准器需要提供一批代表性样本，TensorRT会逐层统计激活值分布，用熵最小化方法确定量化scale。校准集的质量直接决定INT8精度：样本要覆盖真实场景的多样性，数量通常500到1000张即可，过多不会显著提升效果，过少则容易过拟合。校准完成后，scale因子会被固化进plan文件，运行时不再需要校准数据。

动态形状与INT8校准可以同时使用，但要注意：校准过程默认在opt形状下进行，如果实际推理经常落在min或max附近，建议把opt设置为最常用的形状，而不是简单取中间值。

plan文件不可跨TensorRT大版本复用。升级JetPack后，即使GPU型号不变，也应当重新构建引擎，否则可能出现反序列化失败或推理结果异常。

把以上环节串起来看，TensorRT的核心价值在于“用构建时间换运行时间”。Builder阶段做的融合、调优和量化，都是为了让Runtime阶段尽可能少地访问显存、尽可能多地命中Tensor Core。理解这条主线，后续做性能剖析和精度调优时就能快速定位问题出在解析、构建还是执行环节。