模型移植踩坑！PyTorch 到边缘盒子 RKNN 完整部署方案

适用：瑞芯微 RK3588/RK3568 等边缘盒子，RKNN‑Toolkit2，PyTorch 训练模型，覆盖转换、踩坑、推理、性能调优、实际部署整套流程。

整体流程总览

PyTorch训练模型 → 导出ONNX → RKNN‑Toolkit2转RKNN模型 → 模型下发边缘盒子 → RKNN C‑API/Python API推理 → 业务封装上线

重点坑点集中在：ONNX 导出、算子不兼容、精度掉点、输入输出维度、NPU 内存、预处理前后不一致。

一、PyTorch 导出 ONNX（最容易埋坑）

导出代码示例

import torch

model.ｅｖａｌ()

dummy_input = torch.randn(1, 3, 640, 640) # NCHW，和推理输入严格对齐

torch.onnx.export(

    model,

    dummy_input,

    "model.onnx",

    opset_version=17,      # RKNN推荐opset12‑17，不要用过高

    input_names=["images"],

    output_names=["output0"],

    do_constant_folding=True,

    dynamic_axes=None      # 边缘盒子尽量关闭动态shape！动态极易出bug

)

✅导出必做检查

1. 模型必须ｅｖａｌ()模式，关闭 Dropout、BN 训练模式

2. 优先固定输入 shape，动态 shape 能不用就不用，RKNN 对动态支持有限

3. 不要使用torch.nn.functional.grid_sample部分老版本 RKNN 支持差

4. 导出后用 onnxsim model.onnx model_sim.onnx做简化，消除冗余算子

pip install onnxsim

onnxsim model.onnx model_sim.onnx

踩坑 1：直接导出原始 onnx 不做 sim，存在大量 Identity、Cast 冗余算子，转换报错、推理速度暴跌。

二、RKNN‑Toolkit2 模型转换 PC 端（Windows/Linux）

注意：RKNN‑Toolkit2 是PC 端转换工具，边缘盒子只跑推理，不在板子上转模型。

转换脚本模板

from rknn.api import RKNN

rknn = RKNN(verbose=True)

# 配置

rknn.config(

    target_platform='rk3588', # rk3568/rk3566修改这里

    quantized_algorithm='normal',

    quantized_input_type='uint8', # 输入uint8，板子直接传图片，省归一化CPU开销

    quantized_output_type='float32',

    mean_values=[[0,0,0]],

    std_values=[[255,255,255]]

)

# 加载简化后的onnx

rknn.load_onnx(model='model_sim.onnx')

# 构建模型

rknn.build(do_quantization=True, dataset='dataset.txt') # int8量化；不需要量化设False

rknn.export_rknn("model.rknn")

rknn.release()

dataset.txt：量化校准数据集，每行一张图片绝对路径，20‑100 张真实业务图片，不要用随机图

/home/data/img1.jpg

/home/data/img2.jpg

...

高频踩坑清单

1. 算子不支持

报错：Op xxx not support

方案：把不支持算子替换成 RKNN 兼容实现；或者 onnxsim 优化；复杂算子拆分，部分算子用 CPU 回退。

常见坑：aten::unfold、部分自定义算子、torch.where 复杂逻辑。

2. 量化后精度暴跌（最痛苦）

校准数据集必须是真实业务场景图片，不能随机噪声；

尝试quantized_algorithm='kl'；

关键层设置skip_quant跳过量化；

必要时部分分支使用 fp16，牺牲算力换精度。

3. mean/std 配置错误，推理结果完全错乱

重点：如果 rknn.config 填写 mean/std，板子端就不要再做归一化，RKNN 内部完成预处理。两边重复归一化输出直接报废。

4. ONNX NCHW vs NHWC

RKNN 内部 NPU 是 NHWC；load_onnx 会自动转换，不要手动乱改维度。

三、PC 端仿真验证（非常重要，不要直接上板子）

转换完成后，PC 仿真跑一遍，对比 PyTorch 输出和 RKNN 输出，看误差。

# 初始化仿真环境

ret = rknn.init_runtime()

# 输入图片预处理，注意：和dataset校准预处理完全一致

outputs = rknn.inference(inputs=[img_data])

# 和pytorch结果做diff，看最大误差

坑：PC 仿真没问题，上板子异常 → 多半是预处理逻辑不一致、数据类型 uint8/float 混用。仿真环境≠板子真实 NPU。

仿真通过，才把model.rknn拷贝到边缘盒子。

四、边缘盒子端部署（RK3588 为例）

边缘盒子系统：Linux，安装 rknn_api 运行库。

两种开发方式：

1. Python API：快速原型验证

2. C API：产品正式发布，性能最优

Python 推理示例（板子本地运行）

from rknn.api import RKNN

import cv2

rknn = RKNN()

rknn.load_rknn('./model.rknn')

# target=NPU 指定NPU运行

rknn.init_runtime(target='rknn')

img = cv2.imread("test.jpg")

img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

img = cv2.resize(img,(640,640))

# 注意：如果config设置了mean/std，这里直接传uint8图片，不要/255归一化

output = rknn.inference(inputs=[img])

# 后处理：nms、解码等

rknn.release()

重点踩坑：预处理不一致 90% 移植失败的根源

PyTorch：BGR→RGB /resize//255 /mean/std 归一化 NCHW

RKNN：要么写在 config 内部；要么代码手动处理。两套逻辑必须完全一模一样。

经常出现：PC 上效果完美，板子识别全错，就是预处理两边逻辑不一样。

性能相关踩坑

1. 内存泄露：循环推理不要反复 load_rknn、init_runtime，初始化一次循环推理，结束 release。

2. NPU 多模型并发：RK3588 NPU 算力共享，多模型同时跑会降速；大模型注意内存占用。

3. 输入图片不要传 float32，uint8 输入可以节省 CPU 带宽。

五、常见疑难问题汇总

六、生产环境建议

1. 开发流程：PyTorch → onnxsim → RKNN 转换 → PC 仿真校验误差 → 板子实测 → 业务封装

2. 正式产品优先 C API，Python 仅调试；

3. 尽量使用 int8 量化；对精度敏感模块，可混合 fp16；

4. 保存每一步输出：onnx、sim‑onnx、rknn，便于回溯定位问题。

七、避坑核心总结

1. onnxsim 必跑，不要拿原始导出 onnx 直接转换；

2. 能不用动态 shape 就不用；

3. 量化校准数据集必须真实业务图像；

4. 预处理逻辑全局唯一，PC 仿真、转换配置、板子推理三方保持统一；

5. PC 仿真只能做参考，最终必须以边缘盒子 NPU 真实运行结果为准。

家具美容培训

家具维修培训