系列导读:上一篇跑通了第一个推理程序,用的是现成的 MobileNetV2 ONNX 模型。实际项目中你会遇到各种框架训练出来的模型——PyTorch、TensorFlow Lite、PaddlePaddle——这篇把三种主流框架的转换流程全部打通,并重点解决"算子不支持"这个最常见的拦路虎。


一、RKNN-Toolkit2 支持的模型格式

先明确 RKNN-Toolkit2 能直接吃哪些格式:

源框架 支持格式 推荐转换路径
PyTorch .pt / .pth PyTorch → ONNX → RKNN ✅
TensorFlow .pb / SavedModel TF → TFLite → RKNN ✅
TFLite .tflite 直接加载 ✅
PaddlePaddle .pdmodel Paddle → ONNX → RKNN ✅
ONNX .onnx 直接加载 ✅(最通用)
Caffe .caffemodel 直接加载(逐渐淘汰)

核心原则:优先走 ONNX 中间格式。 ONNX 是目前支持最广、算子兼容性最好的路径,除非模型有特殊算子在 ONNX 导出时会丢失,否则统一走 源框架 → ONNX → RKNN


二、PyTorch → RKNN(最常用路径)

2.1 导出 ONNX 的关键注意事项

PyTorch 导出 ONNX 看似简单,但有几个细节不注意就会踩坑:

python

# export_to_onnx.py
import torch

def export_model(model, save_path, input_shape=(1, 3, 640, 640)):
    model.eval()
    dummy = torch.randn(*input_shape)

    torch.onnx.export(
        model,
        dummy,
        save_path,
        opset_version=12,          # ⚠️ 推荐 11-13,不要用 17+
        input_names=["images"],
        output_names=["output"],
        dynamic_axes=None,         # ⚠️ 务必关闭动态 shape,边缘端固定尺寸
        do_constant_folding=True,  # 常量折叠,减小模型体积
        verbose=False
    )
    print(f"✅ 导出成功:{save_path}")

⚠️ 三个常见坑:

  1. opset_version 不要超过 13,RKNN-Toolkit2 对高版本 opset 的某些算子支持不完整
  2. dynamic_axes 必须设为 None,动态 shape 会导致转换失败或推理结果错误
  3. 含有 torch.nn.functional.interpolate 的模型导出后要用 onnxsim 简化,否则会产生冗余算子
2.2 ONNX 简化(强烈推荐)

导出 ONNX 后,先用 onnxsim 做图优化,减少冗余节点,提高转换成功率:

bash

pip install onnxsim
python3 -m onnxsim model.onnx model_simplified.onnx

# 验证简化后的模型是否正确
python3 -c "import onnx; onnx.checker.check_model('model_simplified.onnx'); print('OK')"
2.3 转换为 RKNN

python

# convert_pytorch.py
from rknn.api import RKNN

rknn = RKNN(verbose=False)

rknn.config(
    mean_values=[[0, 0, 0]],        # 根据你的模型实际均值填写
    std_values=[[255, 255, 255]],   # 归一化到 [0,1]
    target_platform="rk3588",
    quantized_dtype="asymmetric_quantized-8",
    optimization_level=3
)

rknn.load_onnx(model="model_simplified.onnx")
rknn.build(do_quantization=True, dataset="./dataset.txt")
rknn.export_rknn("model.rknn")
rknn.release()

三、TFLite → RKNN

TFLite 模型可以直接加载,无需经过 ONNX,是 TF 系模型的最优路径:

3.1 TensorFlow → TFLite 导出

python

# export_tflite.py
import tensorflow as tf

# 加载 SavedModel 或 .pb 模型
converter = tf.lite.TFLiteConverter.from_saved_model("saved_model_dir")

# INT8 量化(可选,RKNN 转换时会再次量化)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]

tflite_model = converter.convert()
with open("model.tflite", "wb") as f:
    f.write(tflite_model)
print("✅ TFLite 导出完成")
3.2 TFLite → RKNN 转换

python

# convert_tflite.py
from rknn.api import RKNN

rknn = RKNN(verbose=False)

rknn.config(
    mean_values=[[127.5, 127.5, 127.5]],  # MobileNet 系列常用均值
    std_values=[[127.5, 127.5, 127.5]],
    target_platform="rk3588",
    quantized_dtype="asymmetric_quantized-8",
)

# 直接加载 tflite,指定输入 shape
rknn.load_tflite(
    model="model.tflite",
    # TFLite 模型通常已含 shape 信息,不需要额外指定
)

rknn.build(do_quantization=True, dataset="./dataset.txt")
rknn.export_rknn("model_from_tflite.rknn")
rknn.release()

💡 TFLite 转换的优势:TFLite 格式本身就是为移动端优化过的,算子集合较小,与 RKNN 的兼容性通常优于直接从 TF SavedModel 转换。


四、PaddlePaddle → RKNN

PaddlePaddle 在国内工业场景用得比较多,推荐路径是 Paddle → ONNX → RKNN:

4.1 安装 paddle2onnx

bash

pip install paddle2onnx paddlepaddle
4.2 导出 ONNX

bash

# 命令行方式(最简单)
paddle2onnx \
    --model_dir ./inference_model \
    --model_filename model.pdmodel \
    --params_filename model.pdiparams \
    --save_file model.onnx \
    --opset_version 12 \
    --enable_onnx_checker True

# 验证
python3 -m onnxsim model.onnx model_simplified.onnx
4.3 后续转换

ONNX 导出后,按照第二节的 PyTorch 路径继续即可,完全相同。


五、算子不支持:最常见拦路虎的系统性解法

这是本篇最重要的一节。实际项目中,模型转换失败 80% 的原因都是算子不支持,必须掌握系统性的排查和解决思路。

5.1 如何快速定位不支持的算子

转换失败时,RKNN-Toolkit2 会输出类似这样的错误:

E Unsupported op: NonMaxSuppression
E Unsupported op: ScatterND

也可以主动查询:

python

# 查询模型中所有算子及支持状态
from rknn.api import RKNN
rknn = RKNN()
rknn.load_onnx("model.onnx")

# 列出所有算子
rknn.list_support_info()  # 输出支持的算子列表
rknn.release()
5.2 四种解决策略(按推荐优先级排序)

策略一:切分模型,不支持的算子在 CPU 上跑

这是最常用也最稳妥的方案。把不支持的算子(通常是后处理部分,如 NMS)切分出来,NPU 跑特征提取,CPU 跑后处理:

python

# 以 YOLOv8 为例:只转换 backbone+neck+head 的特征提取部分
# NMS 等后处理在 C++ 里手写,不进 RKNN

rknn.load_onnx(
    model="yolov8.onnx",
    outputs=["output0"]   # 只取特征图输出,不包含后处理节点
)

策略二:在导出时替换不支持的算子

部分算子在 PyTorch 里有等价的可支持替换写法:

python

# ❌ 不推荐:使用 F.interpolate 的 align_corners=True
x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=True)

# ✅ 推荐:align_corners=False,RKNN 支持更好
x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)

# ❌ 不推荐:torch.nn.SiLU(在某些旧版 RKNN 中不支持)
# ✅ 推荐:等价实现
def silu(x):
    return x * torch.sigmoid(x)

策略三:使用自定义算子(Custom Op)

RKNN-Toolkit2 支持注册自定义算子,适合有特殊计算需求的场景:

python

# 注册自定义算子示例
rknn.config(
    custom_string="my_custom_op",
    ...
)

⚠️ 自定义算子开发复杂度较高,非必要不使用,优先考虑策略一和策略二。

策略四:降级模型,换用兼容性更好的架构

如果模型是自研的,可以在设计阶段就选择 RKNN 兼容性好的算子组合:

避免使用 替代方案
Deformable Conv 普通 Conv
NonMaxSuppression(ONNX 版) CPU 手写 NMS
ScatterND 重新设计网络结构
GridSample(部分版本) 升级 RKNN-Toolkit2 到最新版
5.3 算子支持查询速查

bash

# 查看当前 RKNN-Toolkit2 版本支持的完整算子列表
python3 -c "
from rknn.api import RKNN
rknn = RKNN()
rknn.list_support_info()
"

# 或直接查看官方文档
# https://github.com/airockchip/rknn-toolkit2/tree/master/doc

六、转换质量验证:三步检查法

模型转换完成后,不要直接推到板子,先做三步验证:

步骤一:精度对比(PC 端)

python

# accuracy_check.py
from rknn.api import RKNN
import numpy as np

rknn = RKNN()
rknn.load_rknn("model.rknn")
rknn.init_runtime()

# 用同一张图分别做原始模型推理和 RKNN 推理
# 对比 Top-1 结果是否一致,输出概率差异是否在 3% 以内
input_data = np.random.randint(0, 255, (1, 224, 224, 3), dtype=np.uint8)
outputs = rknn.inference(inputs=[input_data])
print("RKNN 输出 shape:", outputs[0].shape)
print("最大值:", outputs[0].max(), "最小值:", outputs[0].min())
rknn.release()
步骤二:精度分析(发现量化误差层)

python

# 开启逐层精度分析(耗时较长,仅调试时使用)
rknn.accuracy_analysis(
    inputs=["./test_input.npy"],
    output_dir="./accuracy_output",
    target="rk3588",          # 连接板子时可指定 target
    device_id=None
)
# 分析结果在 accuracy_output/ 目录,找 cos_similarity < 0.99 的层重点排查
步骤三:性能预估

python

# 不连板子,在 PC 上估算推理耗时
rknn.eval_perf(is_print=True)
# 输出各层耗时分布,提前发现性能瓶颈

七、不同模型的 mean/std 速查表

这是转换时最容易填错的参数,整理常用模型的标准值:

模型系列 mean_values std_values 备注
ImageNet 预训练(PyTorch) [123.675, 116.28, 103.53] [58.395, 57.12, 57.375] RGB 顺序
MobileNet(TFLite) [127.5, 127.5, 127.5] [127.5, 127.5, 127.5] 归一化到 [-1,1]
YOLO 系列 [0, 0, 0] [255, 255, 255] 归一化到 [0,1]
PaddleDetection [123.675, 116.28, 103.53] [58.395, 57.12, 57.375] 同 ImageNet
人脸模型(RetinaFace) [104, 117, 123] [1, 1, 1] BGR 顺序,注意通道

⚠️ mean/std 填错是推理结果完全错误的最常见原因之一。 转换前务必查清楚原始模型的预处理代码,对照填写。


八、完整转换流程检查清单

转换前:
  ☐ 确认模型输入 shape 固定(无动态 shape)
  ☐ opset_version ≤ 13
  ☐ 用 onnxsim 简化 ONNX
  ☐ 准备 100-300 张有代表性的校准图片
  ☐ 查清楚模型的 mean/std 预处理参数

转换中:
  ☐ 遇到不支持算子 → 优先切分模型
  ☐ verbose=True 查看详细日志

转换后:
  ☐ PC 端模拟推理验证输出 shape 正确
  ☐ accuracy_analysis 确认量化误差 < 3%
  ☐ eval_perf 预估耗时
  ☐ 推到板子做最终验证

九、总结与下篇预告

本篇覆盖了三大框架的完整转换路径,以及算子不支持的系统性解法。掌握这些,你可以把任意主流框架训练的模型转换到 RK3588S NPU 上运行。

下一篇(系列第 5 篇)深入讲 INT8 量化实战:为什么量化会掉精度、校准数据集如何选、混合量化怎么用——把量化这件事彻底搞清楚。


本系列文章列表(持续更新)

Logo

免费领 150 小时云算力,进群参与显卡、AI PC 幸运抽奖

更多推荐