3. ONNX模型基础

在进入RKNN工具链之前，必须先透彻理解ONNX本身。ONNX（Open Neural Network Exchange）是一个开放、跨框架的模型交换标准，它定义了一套统一的计算图格式，让PyTorch、TensorFlow、PaddlePaddle等框架训练出的模型能够转换到同一套中间表示上，再分发到不同推理后端。对RKNN部署而言，ONNX几乎是必经之路——Rockchip官方工具链对ONNX的支持最成熟、算子覆盖最全，因此掌握ONNX的结构与获取方法，是后续量化与NPU适配的地基。

3.1 ONNX的核心设计：Protobuf与计算图

ONNX文件本质上是一个使用Protocol Buffers（protobuf）序列化的二进制文件，扩展名通常为.onnx。它描述了一张有向无环图（DAG），图中包含三类核心元素：节点（Node）、张量（Tensor）和图属性（GraphProto）。理解这三者的关系，是读懂任何ONNX模型的第一步。

核心概念速览：

  • ModelProto：最外层容器，包含模型版本、生产者名称、图结构（GraphProto）以及元数据（如Opset版本）。
  • GraphProto：计算图本体，包含所有节点、初始权重（Initializer）、输入输出张量定义。
  • NodeProto：单个算子实例，包含算子类型（op_type）、输入输出张量名、属性（attribute，如卷积的kernel_shape）。
  • ValueInfoProto：描述张量的名称、数据类型和形状（shape），用于定义图的输入、输出和中间张量。
  • Initializer：直接嵌入图内的权重张量，不需要外部文件加载。

下图展示了ONNX模型从文件到内存解析后的逻辑结构，以及它与RKNN转换器之间的交互关系：

ONNX 文件内部结构 ModelProto（根容器） ir_version / opset_import / producer GraphProto（计算图） NodeProto 列表（算子） • Conv (kernel=3x3, stride=2) • BatchNormalization (eps=1e-5) • Relu • MaxPool (pool_size=2x2) Initializer conv1.weight [64,3,7,7] bn1.bias [64] ValueInfo input: [1,3,224,224] output: [1,1000] AttributeProto（算子属性） Conv: kernel_shape=[3,3], pads=[1,1,1,1] BatchNorm: epsilon=0.00001, momentum=0.9 Transpose: perm=[0,2,3,1] RKNN 转换器视角 步骤1：解析 ModelProto 读取 opset 版本，检查算子兼容性 步骤2：遍历 GraphProto 节点 将 ONNX 算子映射到 RKNN 内部算子 步骤3：权重与形状推导 从 Initializer 加载权重，推断中间张量形状 步骤4：生成 RKNN 计算图 优化算子融合（如 Conv+BN → Conv） 不支持的算子/属性 → 报错：Unsupported Op: xxx → 需修改模型或升级 RKNN 版本 解析

从图中可以看出，RKNN转换器的工作本质上是将ONNX的GraphProto逐节点翻译为RKNN内部计算图。转换失败通常发生在节点映射阶段——当ONNX算子不在RKNN支持列表中，或属性组合不被NPU驱动接受时，就会抛出异常。

3.2 常用算子与NPU适配要点

不同NPU对算子的支持程度差异巨大。RKNN系列NPU对ONNX算子的支持分为完全支持、部分支持（有限制条件）和不支持三类。实际部署中，90%以上的模型由以下基础算子构成：

算子类型 常见用途 RKNN适配注意事项
Conv / ConvTranspose 特征提取、上采样 支持1D/2D/3D；dilation>1时可能降低NPU利用率
BatchNormalization 训练加速、稳定分布 转换时通常被融合进Conv，不单独占用NPU算子
Relu / LeakyRelu / Clip 激活函数 Clip需指定min/max；LeakyRelu的alpha需为常量
MaxPool / AveragePool 下采样 kernel_size和stride需匹配；global pooling需用GlobalAveragePool
Gemm / MatMul 全连接、注意力机制 高维MatMul可能被拆分为多个NPU周期，注意性能
Reshape / Transpose / Concat 张量形状调整 Transpose的perm需在转换时静态确定；Reshape不能有动态维度
Softmax / Sigmoid 分类概率、门控 Softmax的axis属性需明确；大尺度Softmax可能精度下降
Resize / Upsample 图像缩放、特征图放大 需指定mode（nearest/bilinear）和coordinate_transformation_mode

常见坑点：ONNX的Resize算子在opset 11之后属性变化较大，roi和scales可能为空张量。RKNN工具链对某些opset版本的Resize支持不完整，建议在导出ONNX时固定使用opset_version=12，并避免使用antialias等高级属性。

3.3 获取ONNX模型的三种典型方法

获取ONNX模型并非只有“从PyTorch导出”一条路。根据项目来源不同，我们归纳出三种主流方式，每种方式都有其适用场景和注意事项。

方法一：从PyTorch导出（最常用）

PyTorch提供了torch.onnx.export接口，可以将nn.Module转换为ONNX格式。导出过程需要提供一个示例输入张量，用于执行一次前向传播以追踪计算图。关键参数包括opset_version、input_names、output_names以及dynamic_axes（用于动态维度）。

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.onnx",            # 输出文件路径
    opset_version=12,           # 推荐使用12或13
    input_names=["input"],      # 输入张量名
    output_names=["output"],    # 输出张量名
    dynamic_axes={              # 动态维度定义（可选）
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    }
)
print("导出完成")

实操建议：导出前务必调用model.eval()切换到推理模式，否则BatchNorm和Dropout的行为会不一致。如果模型包含torch.jit.ScriptModule或自定义autograd函数，建议先使用torch.jit.trace或torch.jit.script进行封装，再执行导出。

方法二：从TensorFlow / Keras转换

TensorFlow模型不能直接导出为ONNX，需要借助tf2onnx工具。转换时需指定输入节点的名称和形状，TensorFlow的NHWC布局会在转换过程中自动调整为ONNX标准的NCHW布局。

# 安装：pip install tf2onnx
import tensorflow as tf
import tf2onnx

# 加载Keras模型
model = tf.keras.models.load_model("mobilenet_v2.h5")

# 指定输入签名
spec = (tf.TensorSpec((1, 224, 224, 3), tf.float32, name="input"),)

# 转换为ONNX
onnx_model, _ = tf2onnx.convert.from_keras(model, input_signature=spec, opset=13)
# 保存
with open("mobilenet_v2.onnx", "wb") as f:
    f.write(onnx_model.SerializeToString())

方法三：从ONNX Model Zoo或现有模型库获取

如果不想自己训练或转换，可以直接从ONNX Model Zoo（官方模型仓库）或Hugging Face下载现成的ONNX模型。这些模型已经过验证，算子版本兼容性较好，适合作为RKNN转换的测试样例。

# 使用onnxruntime直接加载验证
import onnxruntime as ort
import numpy as np

sess = ort.InferenceSession("resnet18.onnx")
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name

# 构造输入
x = np.random.randn(1, 3, 224, 224).astype(np.float32)
result = sess.run([output_name], {input_name: x})
print("输出形状:", result[0].shape)

3.4 模型验证与常见误区复盘

拿到ONNX文件后，不能直接扔给RKNN转换器，必须先做完整性验证。很多转换失败的问题，根源在于ONNX模型本身存在缺陷或使用了不规范的算子组合。

验证清单（推荐顺序执行）：

  1. 使用onnx.checker.check_model()检查模型结构合法性。
  2. 使用onnx.shape_inference.infer_shapes()推断所有中间张量的形状，确认没有未知维度。
  3. 使用onnxruntime跑一次推理，对比原始框架的输出误差（通常要求cosine相似度 > 0.99）。
  4. 使用netron可视化模型，检查数据流方向、节点连接是否符合预期。
  5. 确认opset版本与RKNN工具链支持范围匹配（建议opset 11~13）。

以下是实际项目中常见的几个误区，值得反复对照自查：

误区 现象 正确做法
导出时忘记model.eval() BN层被当作训练模式，权重被错误折叠 导出前显式调用eval()，并用torch.no_grad()包裹
使用动态维度但未声明dynamic_axes RKNN转换时报shape错误 在导出时明确指定动态轴的符号名称
模型中包含if或loop控制流 ONNX图出现不支持的控制流算子 改用torch.where或torch.max等向量化操作替代
自定义算子未注册 导出时报RuntimeError: Exporting operator xxx not implemented 使用torch.onnx.register_custom_op_symbolic注册映射，或改用标准算子重写
权重精度为float64 RKNN转换时精度不支持 导出前统一转为float32，使用model.float()

特别提醒：ONNX模型中的Constant节点如果包含大型张量（如超过100MB的权重），会导致RKNN转换时内存占用过高。建议在导出前检查模型大小，必要时进行权重剪枝或分块处理。另外，opset_version并非越高越好——RKNN工具链对opset 13+的部分新算子（如Einsum、ScatterND）支持尚不完善，生产环境建议锁定在opset 12。

最后，建议在项目初期就建立一套模型验证流水线：每次从原始框架导出ONNX后，自动执行结构检查、推理比对和RKNN试转换三个步骤。这样可以在进入NPU适配阶段之前，提前暴露80%以上的模型兼容性问题，大幅减少后续排障时间。

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