📚 课程目录第4章模型转换基础

4、模型转换基础：ONNX、PyTorch、TensorFlow模型导出为RKNN格式的流程与工具链

RK3588 的 NPU（神经网络处理单元）原生执行的是 RKNN 格式的模型，而非直接运行 PyTorch 或 TensorFlow 的权重文件。因此，模型转换是部署链路中的第一道关卡，也是决定后续量化精度与推理性能的基石。本章将系统梳理从训练框架到 RKNN 的完整转换流程，重点讲解 ONNX 作为中间表示（IR）的桥接作用、rknn-toolkit2 工具链的安装与调用方式，以及不同框架导出时的关键差异与常见陷阱。

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

RK3588 的 NPU 工具链并不直接解析 PyTorch 的 .pth 文件或 TensorFlow 的 .pb 文件，而是通过 ONNX（Open Neural Network Exchange）作为统一的计算图描述格式。ONNX 本质上是一种协议，它将神经网络描述为有向无环图（DAG），其中节点表示算子（Conv、Relu、Gemm 等），边表示张量数据流。这种设计使得不同训练框架导出的模型可以统一到同一套算子集合上，RKNN 工具链只需针对 ONNX 实现解析器，即可覆盖绝大多数主流模型。

从工程角度看，ONNX 的价值体现在三个层面：算子标准化（不同框架的 Conv 实现统一为 ONNX Conv）、权重布局统一（PyTorch 的 NCHW 与 TensorFlow 的 NHWC 在 ONNX 中明确标注）、调试可视化（Netron 工具可直接查看 ONNX 图结构）。但需要注意，ONNX 并非万能——部分框架特有算子（如 PyTorch 的 F.grid_sample）可能没有 ONNX 对应实现，此时需要手工改写模型或使用自定义算子扩展。

  核心认知：ONNX 是“翻译层”而非“执行层”。它不负责推理计算，只负责描述计算图。RKNN 工具链读取 ONNX 后，会将其转换为 NPU 内部的指令序列，这个过程涉及算子融合（如 Conv+BN+ReLU 融合为一个 NPU 指令）、内存布局转换（NCHW 转 NHWC）以及量化参数嵌入。

4.2 操作流程：PyTorch 模型导出 ONNX 的完整步骤

以 PyTorch 为例，导出 ONNX 的核心 API 是 torch.onnx.export()。以下是一个经过验证的导出模板，适用于大多数视觉模型：

import torch
import torch.nn as nn

# 1. 定义模型并加载权重
model = MyModel()
checkpoint = torch.load('model.pth', map_location='cpu')
model.load_state_dict(checkpoint['state_dict'])
model.eval()

# 2. 构造虚拟输入（必须与训练时的输入尺寸一致）
dummy_input = torch.randn(1, 3, 640, 640)

# 3. 导出 ONNX
torch.onnx.export(
    model,
    dummy_input,
    'model.onnx',
    opset_version=12,          # RKNN 推荐 11~13
    input_names=['input'],     # 输入张量名称
    output_names=['output'],   # 输出张量名称
    dynamic_axes={
        'input': {0: 'batch_size'},   # 动态 batch
        'output': {0: 'batch_size'}
    },
    do_constant_folding=True,  # 常量折叠优化
    verbose=False
)
print("ONNX 导出完成")

导出后必须进行验证，不能直接进入 RKNN 转换。验证方法有两种：静态检查（使用 Netron 打开 ONNX 文件，确认输入输出节点、算子类型、权重形状）和动态验证（使用 onnxruntime 加载模型，对比 PyTorch 输出与 ONNX 输出的数值差异）。

检查项	通过标准	失败处理
算子类型	所有算子均在 RKNN 支持列表内	替换为等效组合（如 HardSwish → ReLU6 + 乘法）
输入输出名称	与后续 RKNN 配置中的名称一致	在导出时显式指定 input_names/output_names
动态轴	RK3588 优先静态 shape	固定 batch=1，分辨率固定为部署值
数值误差	输出差异 < 1e-4（FP32）	检查 BatchNorm 是否冻结、是否有随机层（Dropout）

4.3 工具链实战：rknn-toolkit2 的安装与模型转换

rknn-toolkit2 是瑞芯微官方提供的 Python 包，支持在 x86 Linux 主机上进行模型转换与仿真。安装时需注意 Python 版本（推荐 3.8~3.10）和依赖库（onnx、numpy、opencv-python）。安装完成后，转换流程分为五个步骤：

初始化环境：创建 RKNN 对象并配置目标平台
加载模型：根据原始框架选择 load_onnx / load_pytorch / load_tensorflow
模型构建：执行图优化、算子映射、内存规划
量化准备：加载校准数据集，计算每层动态范围
导出与验证：生成 .rknn 文件，并在 PC 仿真或连接开发板验证

from rknn.api import RKNN

# 1. 创建 RKNN 对象
rknn = RKNN(verbose=True)

# 2. 配置预处理参数（与训练时保持一致）
rknn.config(
    mean_values=[[123.675, 116.28, 103.53]],
    std_values=[[58.395, 57.12, 57.375]],
    target_platform='rk3588',
    quantized_dtype='w8a8'   # 权重 INT8，激活 INT8
)

# 3. 加载 ONNX 模型
ret = rknn.load_onnx(model='model.onnx')
if ret != 0:
    print("加载失败")
    exit(1)

# 4. 构建 RKNN 模型（不量化）
ret = rknn.build(do_quantization=False)
if ret != 0:
    print("构建失败")
    exit(1)

# 5. 导出 RKNN 文件
ret = rknn.export_rknn('model_fp16.rknn')
if ret != 0:
    print("导出失败")
    exit(1)

# 6. 释放资源
rknn.release()

关键风险：rknn.config() 中的 mean_values 和 std_values 必须与训练时的预处理完全一致。如果训练时使用 ImageNet 的 mean=[0.485, 0.456, 0.406] 和 std=[0.229, 0.224, 0.225]，则需转换为 0-255 尺度：mean_values=[[0.485*255, 0.456*255, 0.406*255]]，std_values=[[0.229*255, 0.224*255, 0.225*255]]。否则会导致推理精度严重下降。

4.4 案例拆解：YOLOv8 从 PyTorch 到 RKNN 的完整转换

以 YOLOv8 目标检测模型为例，展示实际转换中遇到的典型问题与解决方案。该模型包含 C2f 模块、SPPF 模块和 Detect 头，其中 C2f 模块使用了 Split 算子，Detect 头包含多分支输出。

初始现象：直接导出 ONNX 后，rknn.load_onnx() 报错“Unsupported operator: Split”。查看 RKNN 算子支持列表，发现 Split 算子仅在特定轴和分段数下受支持。

分析过程：YOLOv8 的 C2f 模块中 Split 将输入张量按通道维度分为两个分支，分段数为 2，轴为 1（NCHW 的 C 维）。RKNN 工具链提示需要将 Split 替换为 Slice 算子组合。

实际操作：修改模型导出脚本，在 torch.onnx.export 之前，将 C2f 模块中的 Split 操作替换为两个 Slice 操作：

# 原始代码
x1, x2 = torch.split(x, split_size_or_sections=[c//2, c//2], dim=1)

# 替换为
x1 = x[:, :c//2, :, :]
x2 = x[:, c//2:, :, :]

结果对比：替换后 ONNX 图中不再出现 Split 算子，RKNN 转换成功。在 RK3588 上使用 NPU 推理，FP16 模型 mAP 与 PyTorch 相差 0.3%，INT8 量化后 mAP 下降 1.2%，在可接受范围内。

复盘经验：遇到算子不支持时，优先查看官方文档的算子支持矩阵；其次考虑用基础算子组合替代；最后才考虑自定义算子（需要编写 C 代码并重新编译 RKNN 运行时）。

4.5 常见误区与总结

在模型转换环节，工程师最容易犯以下错误：

忽略 BatchNorm 融合：PyTorch 的 model.eval() 不会自动融合 BN，需要在导出前手动调用 torch.quantization.fuse_modules() 或使用 torch.jit.script 进行优化。
动态分辨率陷阱：RK3588 NPU 的输入缓冲区是静态分配的，如果模型支持多分辨率输入，必须为每个分辨率单独生成 RKNN 文件，或使用 NPU 的“动态 shape”功能（但会损失部分性能）。
量化校准数据不足：使用 100 张以下图片做校准，可能导致量化尺度偏差。建议使用 500~1000 张覆盖各类场景的图片，且图片需经过与训练相同的预处理。
忽略输出后处理：RKNN 输出的张量布局可能与 PyTorch 不同（如 NHWC vs NCHW），需要在部署代码中做 transpose 或 permute。

效率建议：在团队中建立“模型转换检查清单”，包含算子兼容性检查、数值一致性验证、量化精度评估三个关卡。每次转换新模型时，先跑通 FP16 流程，再尝试 INT8 量化，逐步缩小问题范围。

模型转换是 RK3588 NPU 部署的“第一公里”，也是后续内存优化和性能调优的前提。掌握 ONNX 导出技巧、熟悉 rknn-toolkit2 的配置参数、积累常见算子的替代方案，能够显著缩短从训练到部署的周期。下一阶段，我们将深入 RKNN 模型的内存布局与 NPU 缓存机制，探讨如何通过图优化和内存复用进一步降低推理时的 DDR 带宽压力。