ONNX 和 TensorRT：模型部署加速的 1 分钟速通

你可能遇到过这样的场景：模型在 PyTorch 里训练和验证都没有问题，部署到 Jetson 或服务器 GPU 后却变得很慢；换了推理引擎，输出结果又和原来的结果对不上。

这通常不是模型“突然失效”，而是训练框架、模型交换格式、推理引擎和目标硬件之间还缺少一层清晰的转换与验证。ONNX 和 TensorRT 正好分别解决其中两个关键问题：

• ONNX：用统一的计算图格式描述模型，让模型可以离开训练框架，在不同推理工具之间流转；
• TensorRT：面向 NVIDIA GPU 构建优化后的推理 engine，在构建阶段选择更合适的算子实现、精度和内存规划。

先记住一个容易混淆、但非常重要的关系：

ONNX          = 模型交换格式
ONNX Runtime  = 可以执行 ONNX 的推理运行时
TensorRT      = NVIDIA GPU 推理优化 SDK / 运行时

ONNX 文件本身不会“自动运行”，.engine 也不是跨硬件通用的模型文件。真正的部署结果，取决于模型、输入输出契约、推理运行时和目标硬件是否匹配。

一、ONNX：把模型从训练框架里带出来

ONNX（Open Neural Network Exchange）是一种开放的神经网络交换格式。它用计算图、节点、张量、权重、属性和 opset 版本描述模型的推理过程。PyTorch、TensorFlow 等训练框架可以将模型导出为 ONNX，之后再交给 ONNX Runtime、TensorRT 或其他支持 ONNX 的后端执行。

ONNX 的价值不在于它本身就是最快的运行时，而在于它提供了一层相对稳定的中间表示：训练代码可以继续留在 PyTorch 中，部署侧则可以根据硬件选择不同的执行后端。官方文档对 ONNX 的节点、张量、序列化和 opset 都有更完整的说明，可参考 ONNX Introduction。

为什么图执行通常比 eager 逐算子执行更适合部署

在 PyTorch eager 模式下，前向过程通常表现为一系列算子调用。即使使用了 torch.no_grad() 或 torch.inference_mode()，运行时仍需要处理算子调度、中间张量和内核启动。

导出为静态或半静态计算图后，推理运行时可以提前看到更完整的计算过程，从而进行：

• 常量折叠：把只依赖常量的计算提前完成；
• 冗余节点删除：移除不会影响结果的计算；
• 算子融合：例如把卷积、偏置和归一化相关操作合并；
• 内存复用：让生命周期不重叠的中间张量复用显存；
• 布局与执行计划优化：减少不必要的数据搬运和调度。

ONNX Runtime 的图优化文档 将这些优化分为基础优化、扩展优化和布局优化，并说明了在线优化与离线保存优化图的区别。

用一个简化的 Conv → BatchNorm → ReLU 例子来理解：如果三个算子分开执行，中间结果需要反复写入和读取显存；如果后端能够安全地融合它们，数据搬运和内核启动次数就可能减少。这个例子用于解释优化原理，并不代表任何具体 GPU 上的固定延迟。

因此，“ONNX 一定比 PyTorch 快”并不严谨。速度还取决于：

• 使用的是哪个执行后端，例如 CPU、CUDA、TensorRT 或其他 Execution Provider；
• 模型中的算子是否被完整支持；
• 输入形状是否适合后端优化；
• 是否发生了隐式的数据布局转换或 CPU/GPU 拷贝；
• 是否把预处理、后处理和数据传输也计入了端到端延迟。

二、TensorRT：把 ONNX 变成面向 NVIDIA GPU 的 engine

TensorRT 是 NVIDIA 面向深度学习推理的优化 SDK。典型流程是：把训练框架导出的 ONNX 交给 TensorRT 的 ONNX parser，再由 builder 根据目标 GPU、输入形状和精度配置构建一个序列化 engine。NVIDIA 的 ONNX 部署示例 也是沿着“导出 ONNX → 选择精度 → 构建 engine → runtime 执行”这条路径展开的。

TensorRT 构建 engine 时可能完成以下工作：

• 层融合与垂直融合；
• 在候选 CUDA kernel 中选择更适合目标 GPU 的 tactic；
• 规划中间张量的显存生命周期；
• 在 FP32、FP16、INT8 等精度之间做取舍；
• 为动态输入形状建立 optimization profile；
• 将构建结果序列化为 .engine 文件，供部署侧加载。

这也带来一个工程上的重要限制：engine 不是普通的跨平台模型格式。它通常和 GPU 架构、TensorRT/CUDA 版本、构建配置以及输入 profile 有关。一个环境里构建成功的 engine，不能默认在另一块 GPU 或另一套运行时中直接复用。部署时应记录模型版本、ONNX opset、TensorRT 版本、CUDA 版本、GPU 型号和构建参数。

三、推荐工作流：先保证结果正确，再追求速度

一条可复用的部署流水线是：

PyTorch 训练与验证
    ↓
导出 ONNX
    ↓
ONNX checker + ONNX Runtime 数值对齐
    ↓
图简化或常量折叠（可选）
    ↓
trtexec 构建 TensorRT engine
    ↓
TensorRT 数值对齐与任务指标验证
    ↓
目标 GPU / Jetson 真机性能测试

1. 从 PyTorch 导出 ONNX

下面的代码使用 dynamic_axes 演示 batch 维度动态化。不同 PyTorch 版本的 ONNX 导出器参数会变化；新项目可以根据版本文档优先考虑 dynamic_shapes，具体参数以 PyTorch ONNX 导出文档 为准。

import torch

model.eval()
example_input = torch.randn(1, 3, 640, 640)

torch.onnx.export(
    model,
    (example_input,),
    "model.onnx",
    input_names=["images"],
    output_names=["output"],
    opset_version=18,
    dynamic_axes={
        "images": {0: "batch"},
        "output": {0: "batch"},
    },
)

导出时至少要明确这些信息：输入输出名称、张量布局、数据类型、输入尺寸、动态维度和 opset 版本。检测模型还要把后处理契约写清楚，例如输出是原始 head、解码后的框，还是已经做过 NMS 的结果。

2. 检查 ONNX，并和 PyTorch 做数值对齐

“文件能生成”不代表“模型可用”。先做结构检查，再用同一份输入分别跑 PyTorch 和 ONNX Runtime：

import numpy as np
import onnx
import onnxruntime as ort

onnx.checker.check_model("model.onnx")

session = ort.InferenceSession(
    "model.onnx",
    providers=["CPUExecutionProvider"],
)

torch_output = model(example_input).detach().cpu().numpy()
ort_output = session.run(
    ["output"],
    {"images": example_input.cpu().numpy()},
)[0]

np.testing.assert_allclose(
    torch_output,
    ort_output,
    rtol=1e-3,
    atol=1e-4,
)

对于检测、分割或生成模型，不要只比较最终类别。应同时检查输出 shape、关键中间值、框坐标、置信度、mask 或 token 结果。FP32、FP16 和 INT8 的误差阈值也不应使用同一套固定数字，而应结合模型任务和验证集指标设定。

3. 用 trtexec 构建和测量 TensorRT engine

trtexec 可以在不编写应用代码的情况下构建 engine、运行基准测试和生成 timing cache。NVIDIA 的安装文档特别提醒：pip 安装的 Python 包不包含 trtexec，命令行工作流需要使用带有该工具的安装方式。

固定输入形状可以先从最小命令开始：

trtexec \
  --onnx=model.onnx \
  --saveEngine=model.fp16.engine \
  --fp16

如果 ONNX 使用了动态形状，需要同时指定最小、最优和最大形状：

trtexec \
  --onnx=model.onnx \
  --saveEngine=model.dynamic.fp16.engine \
  --fp16 \
  --minShapes=images:1x3x320x320 \
  --optShapes=images:1x3x640x640 \
  --maxShapes=images:4x3x1280x1280 \
  --shapes=images:1x3x640x640

这里的 --shapes 是本次测试使用的实际输入，--optShapes 则是 builder 用来重点选择 tactic 的形状。二者经常被混为一谈。

四、精度选择：FP32、FP16 还是 INT8

精度选择不是“越低越好”，而是速度、显存、功耗和准确率之间的工程权衡。

FP32：先建立正确性基线

FP32 通常最接近训练和验证时的数值表示，适合先确认模型转换和前后处理没有问题。性能优化最好从一个可复现的 FP32 基线开始，否则后面发现精度变化时，很难判断问题来自 ONNX 导出、TensorRT 构建还是低精度计算。

FP16：多数 NVIDIA GPU 上的第一步优化

FP16 可以减少权重和中间激活的存储压力，并在支持 Tensor Core 的 GPU 上获得更高吞吐。实际收益取决于模型结构和硬件，不能只看模型文件大小。

推荐做法是：

1. 用 FP32 ONNX 和 TensorRT engine 完成数值、任务指标对齐；
2. 构建 FP16 engine；
3. 在同一验证集上比较误差、准确率和端到端延迟；
4. 对异常层或溢出问题保留更高精度。

INT8：需要校准和回归验证

INT8 进一步降低了数据表示的位宽，通常可以减少内存带宽和计算开销，但必须提供具有代表性的校准数据，并重新检查任务级指标。校准数据不需要覆盖整个训练集，但应覆盖真实业务中常见的亮度、尺度、背景、目标大小和输入分布。

如果训练后量化（PTQ）导致精度下降，可以考虑：

• 提高校准数据的代表性；
• 对敏感层保留 FP16 或 FP32；
• 调整 Q/DQ 节点的插入位置；
• 使用量化感知训练（QAT）让模型在训练阶段适应量化误差。

低精度优化完成后，至少要重新跑一次离线验证集，并在目标设备上检查长时间运行的稳定性。

五、动态形状：min、opt、max 不是随便填的

动态形状允许同一个 engine 在一定范围内接受不同的 batch、图像高度或宽度。但动态形状并不意味着“所有输入都同样快”。TensorRT 需要在构建阶段通过 optimization profile 知道允许的范围，并在 opt 形状附近选择更合适的实现。

官方文档将流程概括为：用 -1 标记运行时维度，构建时创建至少一个 optimization profile，运行时选择覆盖实际输入的 profile 并设置输入形状。详见 TensorRT 动态形状基础。

实操时有三个原则：

• min 要覆盖业务允许的最小输入；
• max 要覆盖业务允许的最大输入，但不要为了“保险”无限放大；
• opt 应该选择业务中最常见、最值得优化的输入，而不是简单取平均值。

例如，检测服务大部分请求都是 640×640，偶尔才是 1280×1280，那么 opt 更适合放在 640×640。如果业务只有几种固定尺寸，分别构建多个固定 shape engine，往往比一个范围很宽的动态 engine 更容易获得稳定性能。

六、遇到不支持的算子怎么办

部署报错通常集中在三类问题：opset 版本、算子支持情况和动态 shape。可以按下面的顺序处理：

1. 查阅当前 TensorRT 版本的 ONNX opset 支持指南，确认 parser 是否支持对应算子；
2. 固定输入形状，排除 shape 推导问题；
3. 在训练框架层把复杂或自定义操作改写成更基础的标准算子；
4. 尝试使用合适的 opset 版本重新导出；
5. 使用 Polygraphy 做常量折叠、子图提取和多后端对比；
6. 最后才考虑编写 TensorRT plugin，或更换推理运行时。

不要把“降低 opset”当成万能修复。较低版本可能避开某些 parser 问题，但也可能丢失新算子表达能力。每次更改 opset 或替换算子后，都要重新做数值对齐。

七、结果对不上时，优先排查前后处理

模型转换成功，但结果不对，是部署中最常见的故障之一。建议按照“模型内部 → 输入 → 输出”的顺序排查：

先用同一份随机张量做纯数值对齐。如果 PyTorch、ONNX Runtime 和 TensorRT 对同一输入的输出已经接近，再去查图像预处理、后处理和业务逻辑。这样可以把“图翻译错误”和“业务链路不一致”分开定位。

八、性能测试：不要只看一次平均延迟

真正的部署性能应在目标 GPU 或 Jetson 上测量，并明确 benchmark 的边界：是只测 engine，还是包括图像解码、预处理、数据拷贝和后处理？这两者的数字不能直接比较。

至少记录以下指标：

• 首次推理时间和预热后的稳定延迟；
• P50、P95 或 P99 延迟，而不只是平均值；
• 单路吞吐和多路并发吞吐；
• 峰值显存、主机内存和 engine 文件大小；
• 预处理、推理、后处理分别占用的时间；
• 连续运行数分钟后的温度、功耗和降频情况。

trtexec 很适合建立 TensorRT engine 的基准，但它不能替代完整应用的端到端测试。NVIDIA 的性能基准文档也建议在动态形状场景中明确 profile 和实际测试 shape。

九、ONNX 与 TensorRT 怎么选

一个实用的决策顺序是：

先确定目标硬件与延迟目标
    ↓
再确定输入输出契约与算子覆盖
    ↓
选择 ONNX Runtime、TensorRT 或其他后端
    ↓
最后决定 FP32 / FP16 / INT8 和动态 shape 策略

不要因为“大家都在用 TensorRT”就把所有模型都转成 TensorRT，也不要因为手上已经有一个 ONNX 文件，就默认它在任何硬件上都能直接达到理想速度。部署选型的终点是可验证的端到端指标，而不是某个文件后缀。

十、最后的部署检查清单

在把模型交给应用或服务前，可以快速检查：

• ONNX 结构通过 onnx.checker；
• PyTorch、ONNX Runtime 和 TensorRT 对同一输入的结果完成对齐；
• 输入颜色、布局、归一化和尺寸处理已经记录；
• 输出解码、阈值、NMS 和标签顺序已经记录；
• opset、TensorRT、CUDA、驱动和 GPU 型号已归档；
• 动态 shape 的 min / opt / max 覆盖真实业务请求；
• FP16 或 INT8 的任务级精度完成回归；
• benchmark 明确了 engine-only 还是端到端；
• 真机或目标服务器完成了长时间稳定性测试；
• engine、校准数据、导出命令和验证结果可以复现。

总结起来：ONNX 解决“模型如何从训练框架流向推理后端”，TensorRT 解决“模型如何在 NVIDIA GPU 上高效执行”，而前后处理、数值验证和真机测试决定了这条链路是否真的可用。

参考资料

• ONNX Introduction
• PyTorch ONNX Export
• ONNX Runtime Graph Optimizations
• NVIDIA TensorRT Documentation
• Example Deployment Using ONNX
• TensorRT Dynamic Shapes: Core Concepts
• TensorRT ONNX Opset Guide
• Performance Benchmarking using trtexec