【RK3588S 嵌入式AI系列④】模型转换全流程:PyTorch/TFLite/PaddlePaddle → RKNN 实战指南
系列导读:上一篇跑通了第一个推理程序,用的是现成的 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}")
⚠️ 三个常见坑:
opset_version不要超过 13,RKNN-Toolkit2 对高版本 opset 的某些算子支持不完整dynamic_axes必须设为None,动态 shape 会导致转换失败或推理结果错误- 含有
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 量化实战:为什么量化会掉精度、校准数据集如何选、混合量化怎么用——把量化这件事彻底搞清楚。
本系列文章列表(持续更新)
- ✅ 第1篇:硬件全解析:NPU/CPU/GPU架构与芯片选型指南
- ✅ 第2篇:Linux开发环境从零搭建
- ✅ 第3篇:RKNN SDK快速上手
- ✅ 第4篇:模型转换全流程(本文)
- 🔜 第5篇:INT8量化实战
- … 共16篇
更多推荐




所有评论(0)