3、PyTorch导出ONNX基础：torch.onnx.export核心参数、导出流程、静态shape导出示例

在深度学习模型部署链路中，将PyTorch训练好的模型转换为ONNX（Open Neural Network Exchange）格式是连接训练与推理的关键桥梁。ONNX作为开放的模型交换标准，能够将模型从PyTorch的动态图机制中剥离，转换为静态计算图，从而被TensorRT、OpenVINO、ONNX Runtime等推理引擎高效加载。本章将聚焦于torch.onnx.export的核心参数、标准导出流程以及静态shape导出示例，帮助读者建立对ONNX导出的系统性认知，为后续处理动态shape、算子兼容性等高级主题打下坚实基础。

3.1 torch.onnx.export核心参数详解

torch.onnx.export是PyTorch官方提供的ONNX导出接口，其函数签名中包含了十余个参数，但日常使用中真正需要深入理解的核心参数集中在六个：model、args、f、input_names、output_names、dynamic_axes以及opset_version。理解这些参数的作用机制，是掌握ONNX导出的第一步。

其中，args参数是导出过程的"触发器"。PyTorch通过执行一次前向传播来追踪模型的计算图，因此args的shape直接决定了导出后ONNX模型的输入shape。在静态shape导出场景下，args的每个维度都会被硬编码到ONNX模型中，后续推理时输入必须严格匹配该shape。而opset_version则决定了ONNX算子集的版本，不同版本支持的算子数量和特性不同，例如opset 11引入了更多的动态shape支持能力，opset 13则完善了ReduceSum等算子的定义。

3.2 标准导出流程：从PyTorch到ONNX的五步法

一个规范的ONNX导出流程应当包含五个关键步骤：环境准备、模型预处理、执行导出、结构验证、精度校验。每一步都有其特定的目的和常见陷阱，下面逐一展开。

1. 环境准备：确认PyTorch版本与ONNX版本兼容。建议使用PyTorch ≥1.8.0，ONNX ≥1.9.0。可通过pip list | grep onnx检查已安装版本。
2. 模型预处理：加载模型权重，调用model.eval()切换到推理模式。如果模型包含自定义算子，需要提前注册其符号函数（symbolic function）。
3. 执行导出：构造示例输入，调用torch.onnx.export(model, args, "model.onnx", input_names=["input"], output_names=["output"], opset_version=11)。导出过程中PyTorch会打印导出日志，包含算子转换信息。
4. 结构验证：使用onnx.checker.check_model("model.onnx")检查模型结构合法性，或使用onnx.shape_inference.infer_shapes进行shape推断。
5. 精度校验：使用ONNX Runtime加载导出的模型，对相同的输入分别运行PyTorch模型和ONNX模型，比较输出差异。通常要求最大绝对误差小于1e-4。

# 标准导出流程示例代码
import torch
import torch.nn as nn
import onnx

class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.relu = nn.ReLU()
        self.pool = nn.MaxPool2d(2)
        self.fc = nn.Linear(16 * 112 * 112, 10)
    
    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

# 1. 环境准备与模型预处理
model = SimpleCNN()
model.eval()
dummy_input = torch.randn(1, 3, 224, 224)

# 2. 执行导出（静态shape）
torch.onnx.export(
    model,
    dummy_input,
    "simple_cnn.onnx",
    input_names=["input"],
    output_names=["output"],
    opset_version=11,
    dynamic_axes=None  # 显式指定静态导出
)

# 3. 结构验证
onnx_model = onnx.load("simple_cnn.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX模型结构验证通过")

3.3 静态shape导出实战：以ResNet-18为例

静态shape导出是最简单、最稳定的导出方式，适用于输入尺寸固定的场景，如工业质检中固定分辨率的图像分类。下面以ResNet-18在ImageNet数据集上的分类任务为例，演示完整的静态导出流程，并对比导出前后的推理性能。

场景背景：某安防厂商需要在边缘设备（Jetson Xavier NX）上部署一个图像分类模型，输入为摄像头采集的1080p图像，预处理后固定缩放到224×224。由于输入尺寸恒定，选择静态shape导出以获取最佳推理性能。

初始现象：直接使用PyTorch模型在Jetson上推理，帧率仅达到12 FPS，无法满足25 FPS的实时性要求。团队决定将模型转换为ONNX，并使用TensorRT加速。

关键操作步骤：

1. 加载torchvision预训练的ResNet-18模型，并替换最后的全连接层以适配10类自定义数据集。
2. 构造dummy_input = torch.randn(1, 3, 224, 224)，确保batch size为1，通道数为3。
3. 执行导出，设置opset_version=11，dynamic_axes=None。
4. 使用ONNX Runtime CPU执行器验证输出与PyTorch原始输出的余弦相似度，结果为0.99998。
5. 将ONNX模型导入TensorRT，生成TensorRT engine，batch size固定为1。

结果对比与复盘：通过静态shape导出并配合TensorRT，推理速度提升了约10倍，完全满足实时性要求。复盘发现，静态shape导出的关键优势在于：所有张量的shape在图中完全固定，TensorRT可以针对固定shape进行极致的内核优化和内存布局规划。但代价是灵活性丧失——如果后续需要处理不同分辨率的输入，必须重新导出模型。

3.4 常见误区与总结复盘

在静态shape导出的实践中，初学者常犯的错误集中在以下几个方面，这里逐一剖析并给出规避建议。

• 误区一：忘记切换eval模式。如果模型包含BatchNorm层，在train模式下导出会将当前batch的统计量写入模型，导致推理时归一化错误。正确做法是导出前调用model.eval()。
• 误区二：使用带有梯度的输入。示例输入必须使用torch.randn等普通张量，不能是nn.Parameter或requires_grad=True的张量，否则导出图会包含梯度计算节点。
• 误区三：忽略opset_version的影响。某些PyTorch算子（如torch.repeat_interleave）在低版本opset中不支持，需要升级opset或手动实现替代方案。建议统一使用opset 11以上。
• 误区四：导出后不进行精度校验。ONNX导出过程中可能发生算子融合或精度调整，直接部署可能导致精度下降。必须使用ONNX Runtime对比输出。