4. ONNX作为中间格式

在Jetson平台部署PyTorch模型时，ONNX（Open Neural Network Exchange）扮演着“通用语言”的角色。PyTorch模型无法被TensorRT直接解析，而ONNX作为中间表示（IR），能够将PyTorch的动态计算图“冻结”为静态图，并保留算子级语义，从而让TensorRT的解析器（onnx-tensorrt-parser）能够读取并优化。本章将系统讲解从PyTorch导出ONNX的完整流程、算子兼容性排查方法、动态轴（Dynamic Axes）的配置策略，以及针对Jetson平台的优化技巧。

PyTorch → ONNX → TensorRT 导出链路与关键检查点 PyTorch 模型 .pth / .pt 动态计算图 torch.onnx.export opset_version dynamic_axes input_names / output_names ONNX 模型 静态图 + 动态轴 .onnx TensorRT Engine .engine / .trt FP16/INT8 优化 关键检查点与优化动作 ① 算子兼容性 - 使用 onnx.checker 验证 - 对比 PyTorch 与 ONNX 算子集 - 处理不支持的算子（重写/分解） ② 动态轴设置 - 指定 batch / 宽高为动态 - 设置 dynamic_axes 字典 - 验证动态 shape 推理 ③ 优化与验证 - 简化图结构（onnx-simplifier） - 设置动态 shape 范围 - 对比输出误差（cosine 相似度） 推荐流程： 导出 ONNX 检查算子 设置动态轴 简化与验证 TensorRT 转换

4.1 核心原理：为什么需要ONNX作为中间层

PyTorch采用动态图机制（Define-by-Run），在前向传播过程中实时构建计算图。这种灵活性在训练阶段极具优势，但在部署阶段却成为障碍——TensorRT需要静态或半静态的网络结构以进行层融合、内存复用和内核自动调优。ONNX作为一种开放格式，将PyTorch模型“导出”为静态计算图，同时保留必要的动态维度信息。其核心价值体现在三个层面：格式中立性（ONNX由微软、Facebook等联合维护，不受单一框架锁定）、算子标准化（定义了统一的算子集，便于不同推理引擎解析）、图优化入口（在ONNX阶段可进行常量折叠、冗余节点消除等预处理）。

在Jetson平台上，ONNX还承担着“适配层”的角色。Jetson的TensorRT版本与桌面GPU不同，其支持的算子集和优化策略存在差异。通过ONNX，开发者可以在x86主机上完成模型转换和初步验证，再将优化后的ONNX文件部署到Jetson设备上，实现开发与部署环境的解耦。

4.2 PyTorch导出ONNX的标准步骤

导出过程的核心是调用torch.onnx.export函数。以下是一个典型的导出流程，以ResNet-18分类模型为例：

import torch
import torchvision.models as models

# 1. 加载模型并设为评估模式
model = models.resnet18(pretrained=True)
model.eval()

# 2. 构造示例输入（与模型输入shape一致）
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 定义输入/输出名称（便于后续动态轴引用）
input_names = ["input"]
output_names = ["output"]

# 4. 定义动态轴（此处将batch维设为动态）
dynamic_axes = {
    "input": {0: "batch_size"},
    "output": {0: "batch_size"}
}

# 5. 执行导出
torch.onnx.export(
    model,
    dummy_input,
    "resnet18.onnx",
    input_names=input_names,
    output_names=output_names,
    dynamic_axes=dynamic_axes,
    opset_version=17,          # 根据TensorRT版本选择
    do_constant_folding=True,  # 常量折叠优化
    export_params=True         # 导出权重参数
)

# 6. 验证ONNX模型
import onnx
onnx_model = onnx.load("resnet18.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX模型验证通过")

导出过程中，opset_version的选择至关重要。Jetson平台上的TensorRT版本决定了其支持的ONNX算子集上限。例如，JetPack 5.1自带的TensorRT 8.5支持ONNX opset 17，而JetPack 6.0的TensorRT 8.6支持opset 19。若opset设置过高，可能导致TensorRT解析失败；设置过低，则可能无法表达某些新算子。建议查阅Jetson设备上TensorRT的Release Notes确认支持范围。

4.3 算子兼容性检查：从PyTorch到ONNX的映射

PyTorch算子与ONNX算子并非一一对应。PyTorch中的某些高级操作（如torch.where、torch.topk）在ONNX中可能映射为多个基础算子，而某些自定义操作（如自定义autograd.Function）则可能完全无法导出。算子兼容性检查应遵循以下流程：

  1. 静态检查：使用onnx.checker.check_model验证ONNX模型的结构合法性。
  2. 动态检查：使用onnxruntime（ORT）进行推理，对比PyTorch输出与ONNX输出的数值差异。
  3. 算子映射表：参考PyTorch官方文档中的“ONNX Export”章节，确认目标算子是否在支持列表中。
  4. 不兼容算子处理：对于不支持的算子，可采用三种策略——重写模型（用基础算子组合替代）、自定义算子注册（通过torch.onnx.register_custom_op_symbolic映射）、在ONNX图中手动修改（使用onnx_graphsurgeon）。

以下表格列出了Jetson部署中常见的算子兼容性情况：

PyTorch算子ONNX映射TensorRT支持注意事项
Conv2d / BatchNorm2dConv / BatchNormalization支持（可融合）建议在PyTorch中提前融合BN，减少转换开销
ReLU / SigmoidRelu / Sigmoid支持无特殊处理
torch.reshape / viewReshape / Flatten支持注意动态shape下的维度推断
torch.topkTopK部分支持Jetson上建议改用ArgMax或手动实现
torch.whereWhere支持确保条件张量的shape广播正确
自定义autograd.Function不支持不支持需重写为ONNX可表达的操作

4.4 动态轴设置：灵活性与性能的平衡

动态轴（Dynamic Axes）允许ONNX模型接受不同shape的输入，这在处理可变分辨率图像或动态batch时非常有用。然而，动态轴并非“免费午餐”——它会让TensorRT在构建引擎时无法完全确定张量维度，从而限制某些优化（如内存池预分配、内核选择）。

设置动态轴时，需要遵循以下原则：

  • 最小化动态维度：只将必要的维度设为动态（如batch维），尽量固定其他维度（如通道数、特征图尺寸）。
  • 指定范围：在TensorRT构建引擎时，通过OptimizationProfile设置动态维度的最小、最优、最大值。例如，batch_size范围设为[1, 4, 8]，TensorRT会针对最优值（4）进行深度优化，同时保证在最小和最大值之间正常工作。
  • 避免动态H/W：对于图像模型，若输入分辨率变化范围过大，建议使用固定分辨率或分档处理（如分别导出224x224和416x416两个引擎）。

以下代码展示了在TensorRT中配置动态shape范围：

import tensorrt as trt

# 创建builder和network
builder = trt.Builder(trt.Logger(trt.Logger.WARNING))
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, trt.Logger(trt.Logger.WARNING))

# 解析ONNX模型
with open("resnet18.onnx", "rb") as f:
    parser.parse(f.read())

# 配置动态shape
profile = builder.create_optimization_profile()
input_tensor = network.get_input(0)
profile.set_shape(input_tensor.name, 
                  min_shape=(1, 3, 224, 224),
                  opt_shape=(4, 3, 224, 224),
                  max_shape=(8, 3, 224, 224))
config = builder.create_builder_config()
config.add_optimization_profile(profile)

# 构建引擎
engine = builder.build_engine(network, config)

4.5 案例拆解：YOLOv5s在Jetson Orin上的ONNX导出与优化

以YOLOv5s目标检测模型为例，展示完整的ONNX导出与优化闭环。该模型包含C3模块、SPP模块和Detect头，其中Detect头涉及torch.cat、torch.view和torch.sigmoid等操作，在导出时需特别注意。

背景与目标：在Jetson Orin Nano上部署YOLOv5s，要求输入分辨率640x640，batch_size=1，FP16推理延迟低于15ms。

初始现象：直接使用官方导出脚本生成ONNX后，TensorRT解析报错“Unsupported operator: NonMaxSuppression”。

分析过程：YOLOv5的Detect头在导出时默认包含NMS操作，而TensorRT的ONNX解析器对NMS的支持有限。通过查看导出日志，发现torchvision.ops.nms被映射为ONNX的NonMaxSuppression算子，但该算子在Jetson的TensorRT版本中未被完整实现。

解决方案：采用“端到端分离”策略——将NMS从模型中剥离，在ONNX中仅保留检测头输出（边界框坐标、置信度、类别概率），NMS在TensorRT后处理阶段用CUDA实现。具体操作：

# 修改YOLOv5的export.py，添加 --no-nms 参数
python export.py --weights yolov5s.pt --include onnx --dynamic --no-nms

结果对比：剥离NMS后，ONNX模型大小从29MB降至14MB，TensorRT引擎构建时间从45秒降至12秒，推理延迟从18ms降至11ms（FP16）。

风险与代价：后处理NMS需要额外编写CUDA代码，增加了开发工作量；但换来了更高的灵活性和更低的延迟。

复盘经验：在导出前，应使用torch.onnx.export的verbose=True参数查看详细导出日志，提前识别不支持的算子。同时，建议在模型设计阶段就考虑部署友好性，将NMS等后处理操作移出模型主干。

4.6 常见误区与总结复盘

误区一：盲目追求最新opset版本。Jetson上的TensorRT可能不支持最新opset，导致解析失败。正确做法是查阅TensorRT Release Notes，选择与设备匹配的opset。

误区二：忽略动态轴的性能代价。动态轴会限制TensorRT的层融合和内存优化。如果业务场景中batch_size固定，应使用静态轴；若需动态，务必设置合理的OptimizationProfile。

经验建议：在导出ONNX后，使用onnx-simplifier进行图简化，可消除冗余的Identity节点和Shape操作，平均减少10%-20%的节点数。同时，使用onnxruntime的get_provider接口验证模型在CPU/GPU上的输出一致性。

总结本章核心要点：ONNX导出是PyTorch到TensorRT的关键桥梁，其质量直接影响最终推理性能。开发者应掌握标准导出流程、算子兼容性排查方法、动态轴配置策略，并结合实际部署场景进行针对性优化。在Jetson平台上，建议始终以“最小动态维度、最大静态优化”为原则，在灵活性和性能之间找到最佳平衡点。

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