📚 课程目录第4章 PyTorch模型导出ONNX

4、PyTorch模型导出ONNX：使用torch.onnx.export导出模型、动态轴设置、算子兼容性检查、常见导出错误与解决

在RKNN部署链路中，PyTorch模型转换为ONNX是承上启下的关键环节。ONNX（Open Neural Network Exchange）作为中间表示，其质量直接决定了后续RKNN Toolkit的转换成功率与NPU推理性能。本章将围绕torch.onnx.export的完整使用流程，深入剖析动态轴配置、算子兼容性检查方法以及高频导出错误的定位与修复策略，帮助读者建立一套可复用的工程化导出方法论。

4.1 核心原理：torch.onnx.export的导出机制与静态图追踪

torch.onnx.export本质上是基于TorchScript追踪（Tracing）机制工作的。当调用该函数时，PyTorch会使用给定的示例输入（dummy input）实际执行一遍模型的前向传播，同时记录所有张量操作，将其转换为ONNX计算图中的节点。这意味着导出过程并非解析模型源码，而是“观察”实际执行路径。因此，模型中的控制流（如if语句、for循环）如果依赖于张量值，则只会被记录为执行时的那一条分支，而非完整的逻辑。

  关键理解：追踪导出是“数据驱动”的。示例输入的shape、dtype、乃至数值范围都会影响导出结果。例如，若模型中对输入张量执行了torch.where(x > 0, a, b)操作，追踪时只会记录x > 0这个比较运算，而不会记录条件分支逻辑。这要求工程师在导出前必须明确模型的输入边界。

导出函数的核心签名如下：

torch.onnx.export(
    model,                      # 要导出的PyTorch模型（nn.Module）
    args,                       # 模型输入的元组，例如 (dummy_input,)
    f,                          # 输出文件路径，如 "model.onnx"
    export_params=True,         # 是否导出权重参数
    opset_version=11,           # ONNX算子集版本，RKNN建议使用11~13
    do_constant_folding=True,   # 是否执行常量折叠优化
    input_names=['input'],      # 输入节点名称
    output_names=['output'],    # 输出节点名称
    dynamic_axes=None           # 动态轴配置字典
)

其中opset_version决定了导出时使用的ONNX算子集版本。版本越高，支持的算子越丰富，但RKNN Toolkit对高版本opset的支持可能滞后。根据瑞芯微官方文档，RKNN Toolkit 2.x建议使用opset 11或12，部分模型在opset 13下也能转换，但需逐一验证。opset版本过低则可能缺少某些关键算子（如Resize的坐标变换模式）。

4.2 操作流程：从静态导出到动态轴配置

完整的导出流程分为四个步骤：模型准备 → 静态导出验证 → 动态轴配置 → 验证与检查。以下通过一个典型的图像分类模型（ResNet-18）演示具体操作。

4.2.1 模型准备与静态导出

首先加载预训练模型并设置为eval()模式，这至关重要——eval()模式会关闭Dropout和BatchNorm的训练行为，确保导出图与推理行为一致。

import torch
import torchvision.models as models

model = models.resnet18(pretrained=True)
model.eval()

# 构造示例输入：batch=1, 3通道, 224x224
dummy_input = torch.randn(1, 3, 224, 224)

# 静态导出：所有维度固定
torch.onnx.export(
    model,
    (dummy_input,),
    "resnet18_static.onnx",
    export_params=True,
    opset_version=11,
    input_names=['input'],
    output_names=['output']
)
print("静态导出完成")

静态导出的ONNX模型输入shape固定为[1, 3, 224, 224]。这在批量推理或固定分辨率场景下没有问题，但若需要支持动态batch或可变输入尺寸，则必须配置动态轴。

4.2.2 动态轴设置方法

动态轴通过dynamic_axes参数指定，其格式为字典：{输入/输出名称: {轴索引: 轴名称}}。以下示例将batch维度和高度、宽度维度设为动态：

dynamic_axes = {
    'input': {0: 'batch_size', 2: 'height', 3: 'width'},
    'output': {0: 'batch_size'}
}

torch.onnx.export(
    model,
    (dummy_input,),
    "resnet18_dynamic.onnx",
    export_params=True,
    opset_version=11,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes=dynamic_axes
)

配置后，ONNX模型输入shape变为['batch_size', 3, 'height', 'width']，其中batch_size、height、width为符号变量。需要注意的是，并非所有算子都支持动态shape。例如，某些reshape操作如果依赖固定的张量形状，动态轴可能导致导出失败或推理错误。因此，动态轴配置后必须进行验证。

4.2.3 验证与检查

导出后，使用onnxruntime进行推理验证是最直接的手段。以下代码检查模型是否可加载，并对比PyTorch与ONNX的输出差异：

import onnxruntime as ort
import numpy as np

# 加载ONNX模型
sess = ort.InferenceSession("resnet18_dynamic.onnx")
input_name = sess.get_inputs()[0].name

# 构造不同batch的输入进行测试
for batch in [1, 4, 8]:
    test_input = np.random.randn(batch, 3, 224, 224).astype(np.float32)
    ort_out = sess.run(None, {input_name: test_input})[0]
    print(f"Batch {batch}: ONNX输出shape {ort_out.shape}")

同时，使用onnx.checker.check_model进行结构完整性校验，并使用Netron可视化工具检查计算图结构是否符合预期。

4.3 算子兼容性检查：从ONNX到RKNN的映射预判

ONNX模型能否成功转换为RKNN，取决于算子是否被RKNN Toolkit支持。瑞芯微提供了rknn-toolkit2的rknn.config()和rknn.load_onnx()接口，但最有效的预检手段是在导出阶段就规避不支持的算子。

ONNX算子	RKNN支持情况	常见替代方案
Conv / BatchNorm / Relu	✅ 原生支持	无需处理
Add / Mul / Concat	✅ 原生支持	无需处理
Resize (双线性/最近邻)	✅ 支持（需指定模式）	使用opset 11的`Resize`算子
Softmax	✅ 支持（axis需明确）	指定`axis=-1`
Gather / Unsqueeze	⚠️ 部分支持	避免动态索引，改用固定索引
Einsum	❌ 不支持	拆解为MatMul + Transpose + Reshape
NonMaxSuppression	❌ 不支持	在NPU外处理（CPU后处理）

检查算子兼容性的具体操作流程如下：

导出ONNX后，使用onnxruntime推理，确认模型本身无逻辑错误。
使用Netron打开模型，逐层检查是否有红色警告标记的算子。
运行rknn.load_onnx()预加载，观察是否抛出“unsupported op”错误。
若发现不支持的算子，回到PyTorch源码中修改实现，用等价的基础算子替换。

经验提示：在模型设计阶段就应避免使用过于复杂的自定义算子。例如，如果需要在模型中实现“通道注意力”，建议使用GlobalAveragePooling + Conv + Sigmoid + Mul的组合，而不是自定义SeBlock，因为后者在RKNN转换时可能因算子融合问题导致失败。

4.4 案例拆解：动态shape模型导出失败与修复

以下是一个真实项目中遇到的案例。某目标检测模型使用torchvision.ops.nms()进行后处理，且输入分辨率不固定（训练时使用了多尺度）。工程师尝试导出动态shape模型，但遇到以下错误：

RuntimeError: Exporting the operator 'nms' to ONNX opset version 11 is not supported.

背景与约束：模型需要在RK3588上部署，输入分辨率范围为320×320至640×640，batch固定为1。NMS操作在PyTorch中属于torchvision.ops扩展，ONNX原生算子集中没有对应实现。

分析判断：NMS属于后处理逻辑，不应包含在推理图中。正确的做法是将NMS从模型中剥离，在NPU推理完成后使用CPU或RKNN的后处理接口实现。

实际操作步骤：

修改模型代码，将forward()中的NMS调用移除，仅保留特征提取和回归头输出。
导出时，输出节点改为['boxes', 'scores']，分别对应检测框坐标和置信度。
在RKNN部署代码中，使用cv2.dnn.NMSBoxes()或自定义C++ NMS实现后处理。

结果对比：修改前导出失败；修改后成功导出，且RKNN转换时间从失败变为约2分钟完成。推理性能方面，由于移除了NMS，NPU计算量减少约5%，整体端到端延迟降低8%。

风险与代价：后处理移至CPU增加了主控负载，但RK3588的CPU性能充裕，实测CPU占用率仅增加3%。

复盘经验：导出前必须梳理模型的计算图边界，将非NPU友好的操作（NMS、排序、动态索引）全部剥离。这不仅是技术问题，更是架构设计问题。

4.5 常见导出错误与解决策略

以下汇总了高频出现的导出错误及其解决方案，这些错误在RKNN部署实践中反复出现。

错误现象	根本原因	解决方案
`RuntimeError: Exporting the operator 'aten::std' to ONNX opset version 11 is not supported.`	模型使用了ONNX不支持的PyTorch算子（如`std`、`var`、`einsum`）	用`sqrt(mean(x^2) - mean(x)^2)`等基础运算替代；或升级opset版本（需验证RKNN兼容性）
`ValueError: Some axes are unused in the model.`	`dynamic_axes`中指定的轴在模型中实际未使用	检查模型输入输出，删除未使用的动态轴定义
`RuntimeError: Input type (double) not supported.`	示例输入使用了`float64`，而模型权重为`float32`	确保`dummy_input`的dtype为`torch.float32`
`onnxruntime.capi.onnxruntime_pybind11_state.InvalidGraph`	ONNX图结构错误，通常由常量折叠或自定义算子引起	设置`do_constant_folding=False`重新导出；检查是否有未注册的自定义算子
RKNN转换时报`Unsupported op: Resize`	opset版本过高，Resize算子的坐标变换模式不被RKNN支持	将`opset_version`降为11；或在PyTorch中使用`F.interpolate`并指定`mode='nearest'`

4.5.1 动态轴配置的常见误区

动态轴并非“越多越好”。在RKNN部署中，动态shape会显著增加转换复杂度，并可能导致NPU内存分配效率下降。建议遵循以下原则：

优先固定batch为1：大多数边缘部署场景batch=1，固定后可简化优化。
高度和宽度尽量固定：如果模型包含全连接层或需要固定尺寸的Reshape，动态分辨率会导致导出失败。
使用opset_version=11作为默认值：opset 11是RKNN支持最成熟的版本，避免使用opset 13+的新特性。

风险警示：当动态轴与torch.onnx.export的do_constant_folding=True结合时，某些shape相关的常量可能被错误折叠，导致动态轴失效。若发现动态shape未生效，请尝试设置do_constant_folding=False。

4.6 总结复盘：导出质量的三重验证

一个合格的ONNX导出必须通过三重验证：结构验证、数值验证、目标平台验证。结构验证使用onnx.checker和Netron；数值验证对比PyTorch与ONNX的输出差异（通常要求余弦相似度>0.999）；目标平台验证则是直接调用rknn.load_onnx()进行试转换。只有三重验证全部通过，才能进入后续的量化与部署阶段。

在实际项目中，建议将导出过程封装为可复用的脚本，并纳入CI/CD流程。每次模型更新后自动执行导出与验证，可有效避免“模型能跑但导不出”的尴尬局面。记住：导出不是终点，而是RKNN部署的起点，一个结构清晰、算子规范的ONNX模型，能为后续的量化精度调优节省大量时间。