3、ONNX模型基础：掌握ONNX的模型结构、节点与张量概念、如何导出或获取ONNX模型（PyTorch/TensorFlow转换）

在瑞芯微NPU（如RK3588、RK3566等）的部署流程中，ONNX（Open Neural Network Exchange）扮演着“通用交换格式”的关键角色。无论是PyTorch训练出的模型，还是TensorFlow产出的模型，最终都需要转换为RKNN格式才能被NPU执行。而ONNX正是这一转换链路的中间桥梁。理解ONNX的底层结构、节点与张量的组织方式，以及如何从主流框架中正确导出ONNX模型，是确保后续RKNN转换顺利、精度不损失的前提。本章将系统拆解ONNX的模型结构、核心概念，并给出从PyTorch与TensorFlow导出ONNX的完整方法论与实战案例。

3.1 ONNX模型结构：计算图、节点与张量

3.1.1 核心原理：ONNX的层次化结构

ONNX模型本质上是一个有向无环图（DAG），它由三个核心层级组成：模型层（ModelProto）、图层（GraphProto）和节点层（NodeProto）。模型层是顶层容器，包含图的定义、算子集版本（opset）、生产者信息和元数据。图层是模型的主体，它定义了计算图的输入（GraphInput）、输出（GraphOutput）以及所有中间节点。节点层则是计算的基本单元，每个节点代表一个算子（如Conv、Relu、Add等），节点通过张量（Tensor）连接形成数据流。

理解这一结构对后续RKNN转换至关重要。RKNN-Toolkit在解析ONNX时，会遍历图中的每个节点，将其映射到NPU支持的算子或组合算子。如果ONNX图中存在NPU不支持的算子（如某些自定义算子），转换就会失败或产生大量子图切分，导致推理性能下降。

3.1.2 节点与张量的关系

在ONNX中，张量（Tensor）是数据的基本载体，它可以是模型的输入、输出，也可以是节点之间的中间结果。每个张量都有名称、数据类型（float32、int8等）和维度信息。而节点（Node）则定义了张量之间的运算关系。例如，一个Conv节点接收输入张量input和权重张量conv1.weight，输出张量conv1.output。节点通过attribute字段存储算子参数（如卷积的stride、padding、dilation等）。

在RKNN转换中，张量的布局（Layout）和数据类型是高频问题。PyTorch默认使用NCHW布局，而部分TensorFlow模型使用NHWC。ONNX标准本身不强制布局，但RKNN-Toolkit在解析时会根据算子类型推断布局。如果布局推断错误，会导致输出结果完全错误。因此，在导出ONNX时，必须明确指定输入张量的布局，并在转换时通过input_shape_list参数进行校正。

3.2 从PyTorch导出ONNX模型

3.2.1 操作流程：torch.onnx.export的核心参数

PyTorch提供了torch.onnx.export函数，用于将nn.Module转换为ONNX格式。其核心参数包括：model（待导出模型）、args（模型输入的示例张量）、f（输出文件路径）、input_names和output_names（指定输入输出张量名）、opset_version（算子集版本）、dynamic_axes（动态轴定义）。

import torch
import torchvision.models as models

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

# 2. 构造示例输入（与训练时的预处理一致）
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 导出ONNX
torch.onnx.export(
    model,                      # 待导出模型
    dummy_input,                # 示例输入
    "resnet18.onnx",            # 输出路径
    export_params=True,         # 导出权重
    opset_version=12,           # 算子集版本，RKNN建议≥11
    do_constant_folding=True,   # 常量折叠优化
    input_names=["input"],      # 输入张量名
    output_names=["output"],    # 输出张量名
    dynamic_axes={              # 动态轴定义（此处固定batch）
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    }
)
print("ONNX导出完成")

3.2.2 案例拆解：ResNet18导出与验证

以一个实际项目为例：我们需要将ResNet18部署到RK3588上。初始现象是，直接使用默认参数导出ONNX后，在RKNN-Toolkit中转换失败，报错Unsupported operator: aten::max_pool2d。分析发现，PyTorch导出的ONNX中，MaxPool算子的属性ceil_mode为False，但RKNN-Toolkit的算子库要求该属性必须显式指定。解决方法是升级opset版本至12，并添加torch.onnx.export的operator_export_type参数为torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK，将不支持的ATen算子转换为ONNX标准算子。

操作步骤如下：

1. 将opset_version从11提升到12，重新导出。
2. 使用onnx.checker.check_model验证模型完整性。
3. 使用onnxruntime进行推理，对比PyTorch输出，确保误差在1e-4以内。
4. 在RKNN-Toolkit中加载ONNX，设置target_platform='rk3588'，成功转换。

结果对比：修正后的ONNX模型在RKNN模拟器上的推理结果与PyTorch输出最大误差为2.3e-5，满足部署要求。复盘经验：导出ONNX后必须进行双重验证——先用onnxruntime验证逻辑正确性，再用RKNN-Toolkit的模拟器验证NPU兼容性。

3.3 从TensorFlow导出ONNX模型

3.3.1 核心原理：tf2onnx的转换机制

TensorFlow模型转换为ONNX通常使用tf2onnx工具。其核心原理是将TensorFlow的GraphDef或Keras模型解析为中间表示，再映射为ONNX算子。与PyTorch不同，TensorFlow的算子粒度更细（如Conv2D、BiasAdd、Relu是分开的），而ONNX中这些通常合并为一个Conv节点（包含bias属性）。因此，tf2onnx需要做算子融合（fusion）操作。

# 方式一：命令行转换（适用于.pb文件）
python -m tf2onnx.convert \
    --graphdef_path model.pb \
    --output model.onnx \
    --inputs input:0 \
    --outputs output:0 \
    --opset 12 \
    --fold_const

# 方式二：Python API（适用于Keras模型）
import tf2onnx
import tensorflow as tf

# 加载Keras模型
model = tf.keras.models.load_model("mobilenet_v2.h5")
# 指定输入签名
spec = (tf.TensorSpec((1, 224, 224, 3), tf.float32, name="input"),)
# 转换
onnx_model, _ = tf2onnx.convert.from_keras(model, input_signature=spec, opset=12)
# 保存
with open("mobilenet_v2.onnx", "wb") as f:
    f.write(onnx_model.SerializeToString())

3.3.2 常见误区与排障

TensorFlow导出ONNX时最常见的误区是输入输出张量名不匹配。在TensorFlow中，张量名通常带有:0后缀（如input:0），而ONNX中不允许冒号。tf2onnx会自动处理，但用户需确保在--inputs参数中正确指定。另一个高频问题是NHWC与NCHW布局混淆。TensorFlow默认使用NHWC，而ONNX标准更倾向于NCHW。tf2onnx会插入Transpose节点来转换布局，但这会增加图复杂度。建议在RKNN转换时，通过input_shape_list参数直接指定NCHW布局，让RKNN-Toolkit自动优化掉多余的Transpose。

3.4 ONNX模型验证与优化

3.4.1 验证流程：从结构到数值

无论从哪个框架导出，ONNX模型都必须经过严格的验证才能进入RKNN转换阶段。验证分为三个层次：结构验证、数值验证和兼容性验证。

• 结构验证：使用onnx.checker.check_model检查模型是否符合ONNX规范，包括节点属性、张量维度、输入输出定义等。
• 数值验证：使用onnxruntime（CPU或CUDA）加载模型，输入随机数据或真实样本，与原始框架的输出进行对比。通常要求最大误差小于1e-4。
• 兼容性验证：使用RKNN-Toolkit的rknn.load_onnx加载模型，查看是否报错或警告。重点关注Unsupported operator和Warning: fallback to CPU等提示。

3.4.2 优化技巧：为NPU部署精简模型

在导出ONNX后，可以使用onnx-simplifier工具进行图优化。该工具会执行常量折叠、冗余节点消除、算子融合等操作，减少模型中的节点数量。例如，将Conv + BatchNorm + Relu三个节点融合为一个Conv节点（将BN参数融合进卷积权重）。这不仅能减小模型体积，还能减少RKNN转换时的算子映射开销。

# 安装onnx-simplifier
pip install onnx-simplifier

# 命令行简化
python -m onnxsim resnet18.onnx resnet18_sim.onnx

# Python API
import onnx
from onnxsim import simplify
model = onnx.load("resnet18.onnx")
model_sim, check = simplify(model)
assert check, "简化失败"
onnx.save(model_sim, "resnet18_sim.onnx")

3.5 本章总结与常见误区

ONNX模型准备是RKNN部署流程中最容易出错的环节之一。总结本章核心要点：ONNX是框架无关的中间表示，其计算图结构决定了NPU的映射效率。在导出时，务必固定输入尺寸、选择合适的opset版本（建议≥11）、验证输出一致性。从PyTorch导出时，注意torch.onnx.export的参数配置；从TensorFlow导出时，关注布局转换和控制流处理。最后，使用onnx-simplifier进行图优化，可以显著提升RKNN转换后的推理性能。

常见误区归纳：

1. 忽略opset版本：低版本opset可能缺少某些算子定义，导致转换失败。
2. 动态shape滥用：RKNN-Toolkit对动态shape支持有限，尽量固定所有维度。
3. 不验证直接转换：跳过onnxruntime验证，可能导致RKNN推理结果错误。
4. 忽略布局差异：PyTorch的NCHW与TensorFlow的NHWC混用，导致输出错乱。

掌握这些基础后，下一步即可进入RKNN-Toolkit的模型转换与量化环节，将ONNX模型高效地映射到瑞芯微NPU上。