4. ONNX模型导出与视觉任务适配
TensorRT 本身并不直接读取 PyTorch 的 .pt 权重，它需要一份中间表示作为输入，而 ONNX 就是当前视觉模型部署链路中最通用的选择。理解 ONNX 导出，本质上要回答三个问题：模型的计算图是否被完整、正确地记录下来；输入输出张量的形状语义是否与推理引擎的预期一致；图中出现的算子是否都能被 TensorRT 解析器识别并映射为高效的 kernel。这三件事任何一件出问题，后面的 INT8 校准、层融合和显存优化都无从谈起。
导出前的模型准备：eval 模式与假输入
PyTorch 的 torch.onnx.export 通过一次前向追踪（trace）来记录计算图，因此模型必须处于推理状态。若忘记调用 model.eval()，BatchNorm 会使用当前 batch 的统计量而非训练时累积的 running_mean/running_var，Dropout 也会随机置零，导出的图会带上训练期行为，导致推理结果与预期不符。同时需要构造一个形状正确的假输入张量，它的 dtype、通道顺序和空间尺寸都会写进 ONNX 的输入签名。
import torch
model = MyDetector(num_classes=80).cuda()
model.load_state_dict(torch.load("detector.pth"))
model.eval()
dummy = torch.randn(1, 3, 640, 640, device="cuda")
torch.onnx.export(
model,
dummy,
"detector.onnx",
input_names=["images"],
output_names=["boxes", "scores", "labels"],
opset_version=13,
dynamic_axes={
"images": {0: "batch", 2: "height", 3: "width"},
"boxes": {0: "batch", 1: "num_det"},
"scores": {0: "batch", 1: "num_det"},
"labels": {0: "batch", 1: "num_det"},
},
do_constant_folding=True,
)
这里 opset_version 的选择需要权衡：opset 11 及以上支持更完整的动态形状语义，opset 13 对 Resize、ScatterND 等算子的定义更贴近视觉模型的实际用法，而 TensorRT 8.x 对 opset 13～17 的支持较为成熟。若模型中含有较新的算子，盲目提高 opset 反而可能让 TensorRT 解析失败，因此建议以 TensorRT 版本文档中标注的支持上限为准。
动态轴设置：让同一份引擎适配多分辨率
视觉任务中，分类模型通常固定 224×224 或 384×384，而检测与分割模型往往需要处理不同尺寸的输入。动态轴的作用就是把某些维度标记为符号维度，使 TensorRT 在构建优化配置文件（Optimization Profile）时能够为一段形状范围分配显存。需要特别注意的是，动态轴只声明“这个维度可以变化”，真正的取值范围要在 TensorRT 侧通过 profile.set_shape() 指定最小、最优、最大三组形状。
对于检测模型，常见做法是只把 batch 和空间维度设为动态，而把通道维固定为 3。若把通道也设为动态，TensorRT 会失去对卷积输入通道的静态推断能力，导致部分融合策略失效。分割模型还要留意输出张量的空间维度是否与输入保持一致的动态关系，否则上采样层在最大形状下可能触发显存溢出。
常见算子兼容性问题与处理策略
导出成功并不等于 TensorRT 能顺利解析。视觉模型中最容易出问题的算子集中在几类：一是 Resize，PyTorch 的 interpolate 在 align_corners、mode 参数不同时会生成不同属性的 ONNX 节点，TensorRT 对 nearest 与 linear 的支持程度不同，遇到不支持的组合需要改写为固定尺寸的 Resize 或替换为反卷积；二是 ScatterND 与 GatherND，常见于分割后处理中的索引操作，TensorRT 8 之前支持有限，可考虑把后处理移出网络，在 CPU 或自定义插件中完成；三是 NonMaxSuppression，虽然 TensorRT 提供了 EfficientNMS_TRT 插件，但需要把原始 ONNX 中的 NMS 节点替换为插件节点，否则解析器会直接报错。
另一个高频问题是动态 shape 与 Reshape 的冲突。当模型中存在 view 或 flatten 操作，且输入空间维度是动态的，ONNX 会生成带 -1 的 shape 张量。TensorRT 在构建时需要推断出具体的维度值，若 -1 无法从上下文中唯一确定，就会在 opt 形状下报维度不匹配。解决办法是尽量让 Reshape 的目标形状只依赖 batch 维，或者把空间维度的动态范围限制在能被整除的取值上。
排查顺序建议：先用 onnx.checker.check_model() 验证图结构合法性，再用 polygraphy surgeon sanitize 做常量折叠与冗余节点清理，最后用 trtexec --onnx=model.onnx --verbose 观察解析器在哪一层失败。日志中 Unsupported operator 或 Cannot infer shape 是最直接的线索。
分类、检测、分割三类模型的适配差异
分类模型结构最简单，通常只有一条主干加一个全连接或池化输出，动态轴一般只需声明 batch 维，导出后几乎不会遇到算子兼容问题。检测模型的关键在于输出头的组织方式：如果模型内部已经包含解码逻辑（把 anchor 偏移还原为绝对坐标），导出时要把这些操作一并 trace 进去，否则 TensorRT 拿到的只是原始偏移量，后处理仍需在外部完成。分割模型则要关注输出分辨率与输入分辨率的比例关系，若网络中有多次下采样再上采样，动态空间维度会让上采样层的输出尺寸计算变得复杂，建议在导出前把上采样改为固定倍率，或使用 scale_factor 而非 size 参数。
| 任务类型 | 典型动态轴 | 高风险算子 | 适配建议 |
|---|---|---|---|
| 分类 | batch | 无 | 固定输入尺寸，直接导出 |
| 检测 | batch, height, width | NonMaxSuppression, ScatterND | NMS 替换为 EfficientNMS_TRT 插件 |
| 分割 | batch, height, width | Resize, Reshape | 上采样用固定倍率，避免 -1 维度 |
导出后的验证闭环
一份可用的 ONNX 模型需要经过数值对齐验证。做法是用同一组输入分别跑 PyTorch 和 ONNX Runtime，比较输出张量的最大绝对误差。对于检测模型，还要比较解码后的框坐标和类别分数；对于分割模型，比较逐像素的 argmax 结果。若误差在 1e-4 量级以内，说明导出过程没有引入结构性偏差，可以进入 TensorRT 构建阶段。若误差偏大，优先检查 eval() 是否生效、是否有自定义算子被 trace 成了 ATen 节点，以及 do_constant_folding 是否意外改变了数值行为。
实用检查清单：① 模型已调用 eval()；② 假输入形状与部署时一致；③ dynamic_axes 覆盖所有需要变化的维度；④ 输出节点名称与后处理代码中的引用一致；⑤ 用 ONNX Runtime 与 PyTorch 做数值对齐；⑥ 用 trtexec 验证解析器能否完整构建引擎。
完成 ONNX 导出与验证后，模型就具备了进入 TensorRT 构建阶段的条件。下一阶段的核心工作是把这份计算图交给解析器，通过优化配置文件限定形状范围，并观察层融合与精度选择带来的实际收益。