文章

RKNN

RKNN,瑞芯微电子(Rockchip),RK3588。

RKNN

简介

RKNN(Rockchip Neural Network)是瑞芯微(Rockchip)专为其NPU(神经网络处理器)开发的深度学习推理引擎。

RKNN是适配该NPU的软件框架,提供模型转换、量化及推理功能,可将TensorFlowPyTorch等模型转换为RK3588等处理器可执行的.rknn格式。

RKNN与RK3588等瑞芯微(Rockchip)推出的高性能处理器,在AI加速和嵌入式推理场景中紧密关联。

RKNN(rknn-toolkit2)下载地址:rknn-toolkit2 GitHub

RKNN包括:

  • RKNN-Toolkit2是一款软件开发工具包,支持用户在PC端和瑞芯微NPU平台上进行模型转换、推理及性能评估。
  • RKNN-Toolkit-Lite2为瑞芯微NPU平台提供Python编程接口,帮助用户部署RKNN模型并加速AI应用的实现。
  • RKNN Runtime为瑞芯微NPU平台提供了C/C++编程接口,旨在帮助用户部署RKNN模型并加速AI应用的实现。
  • RKNPU内核驱动负责与NPU硬件交互,该驱动已开源,可在Rockchip内核代码中找到。
  • RKNN Model Zoo基于 RKNPU SDK 工具链开发, 提供了目前主流算法的部署例程. 例程包含导出RKNN模型, 使用 Python API, CAPI 推理 RKNN 模型的流程.

瑞芯微:

RKNN-Toolkit2 与 RKNN-Toolkit-Lite2

工具对比:

工具功能架构使用场景
RKNN-Toolkit2模型转换、量化、推理x86_64、arm64开发环境(PC端)
RKNN-Toolkit-Lite2模型推理aarch64部署环境(Rockchip设备)

模型转换需要 RKNN-Toolkit2:将 PyTorch 等训练框架的模型转换为 RKNN 格式,需要使用 RKNN-Toolkit2(完整版)。

RKNN-Toolkit-Lite2 的定位:这是一个推理框架,主要用于在 Rockchip 设备上加载和运行已经转换好的 RKNN 模型,而不是用于模型转换。

虽然 RKNN-Toolkit2 主要用于 PC(电脑)端(x86_64),但自2.3.0版本开始,瑞芯微也提供arm64指令架构处理器可使用的.whl轮子。

PC安装RKNN-Toolkit2

在PC上使用RKNN是为了可以在电脑上将Pytorch等模型转换为rknn模型。

rknn-toolkit2 提供了 x86_64arm64 架构CPU可用的Python包(wheel, .whl)。

下载

RKNN的github仓库为airockchip/rknn-toolkit2

若要将模型转为rknn格式,需要安装rknn-toolkit2

下载地址:github.com/airockchip/rknn-toolkit2

可以在TagsReleases中选择不同版本rknn-toolkit2并下载。 目录结构如下,小于2.0.0-beta0的版本可能有所不同:

1
2
3
4
5
6
7
8
9
10
11
12
rknn-toolokit2/
|-autosparsity/
|-doc/
|-res
|-rknn-toolkit-lite2/
|-rknn-toolkit2/
  |-doc/
  |-docker/
  |-examples/
  |-packages/
|-rknpu2/
|-...  # 其他文件,READMD.md等

安装

  • rknn-toolkit2/packages/目录,可以找到类似**linux**_x86_64.whl的Python包,用于安装在x86_64设备,如个人电脑。
  • requirements**.txt文件是rknn-toolkit2的Python依赖。
  • rknn-toolkit2/docker/docker_file/ubuntu_20_04_cp38/目录中有Dockerfile**文件定义的docker镜像。

安装RKNN:

  • 方法1:通过pip install **_x86_64.whl命令安装。
  • 方法2:使用docker镜像。
    构建镜像
    1
    2
    3
    4
    
    # -t 生成的镜像的 名字/标签,Dockerfile在当前目录且文件名是Dockerfile
    docker build -t your-app-image .
    # -f 指定Dockerfile文件路径
    docker build -t your-app-image -f /path/to/Dockerfile
    

    创建Docker容器,这是一个示例需根据情况使用不同的创建命令:

    1
    
    docker run -p 8000:8000 your-app-image
    

我是直接在Python虚拟环境通过pip安装的。

将模型转为rknn格式

PyTorchTensorFlow模型,转换为RK3588等处理器可以运行的.rknn格式,RKNN框架提供了多种方式。 一种是将训练好的模型转换为.onnx格式,然后使用相同的步骤转为.rknn模型;也可以分别采用RKNN中与PyTorch、TensorFlow对应的方法完成转换操作。

  • 将模型转为rknn格式,需要使用rknn-toolkit2
  • RKNN框架支持一维卷积nn.Conv1d

Hailo则不支持nn.Conv1d(截至2025.10.08)。

onnx转rknn

导包:

1
from rknn.api import RKNN

初始化转换器:

1
rknn = RKNN(verbose=True)

配置模型参数:

1
2
3
4
5
6
rknn.config(
    mean_values=[[123.675, 116.28, 103.53]],  # ImageNet标准化参数
    std_values=[[58.395, 58.395, 58.395]],
    target_platform="rk3588",
    quantize_dtype="asymmetric_quantized-8"  # 启用INT8量化
)

其中,mean_valuesstd_values 对转换后的rknn模型的表现(如准确率)有很大影响,这两个参数需要根据训练样本计算得到。

加载ONNX模型:

1
2
3
4
ret = rknn.load_onnx(model="model.onnx")
if ret != 0:
    print("Load ONNX failed!")
    exit(ret)

转换模型:

1
2
3
4
ret = rknn.build(do_quantization=True, dataset="dataset.txt")  # 量化需提供校准数据集
if ret != 0:
    print("Build RKNN failed!")
    exit(ret)

导出RKNN模型:

1
2
3
4
ret = rknn.export_rknn("model.rknn")
if ret != 0:
    print("Export RKNN failed!")
    exit(ret)

释放资源:

1
rknn.release()

其他模型转rknn

对于非onnx模型,只需将load_onnx()方法替换为对应的load_xxx方法,即可实现模型转为rknn格式。

PyTorch模型可使用load_pytorch()方法。

TensorFlow模型可使用load_tensorflow()方法。

RK3588开发板环境配置

开发板系统准备

  1. 刷写系统:使用官方Debian/Ubuntu固件。
  2. 安装依赖:根据自己的项目安装依赖,如Python。

安装RKNN-Toolkit-Lite2

1.创建python虚拟环境(推荐)。

2.安装RKNN-Toolkit-Lite2

  • rknn-toolkit2仓库, 包括rknn-toolkit2 和 rknn-toolkit-lite2。
  • .whl文件路径:rknn-toolkit-lite2/packages/
  • 在虚拟环境中运行pip install xxx_aarch64.whl安装即可。

在RK3588上运行RKNN模型

编写运行脚本

DeepSeek给的示例代码:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
import cv2
import numpy as np
from rknnlite.api import RKNNLite

# 初始化RKNN Lite
rknn_lite = RKNNLite()

# 加载模型
ret = rknn_lite.load_rknn("model.rknn")
if ret != 0:
    print("Load RKNN failed!")
    exit(ret)

# 初始化运行时
ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0)  # 指定NPU核心
if ret != 0:
    print("Init runtime failed!")
    exit(ret)

# 图像预处理
img = cv2.imread("test_image.jpg")
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = cv2.resize(img, (224, 224))
img = np.expand_dims(img, axis=0).astype(np.float32)

# 执行推理
outputs = rknn_lite.inference(inputs=[img])

# 后处理(示例:分类模型)
pred = np.argmax(outputs[0])
print("Predicted class:", pred)

# 释放资源
rknn_lite.release()

根据自己的模型作相应修改。我实际使用一维卷积模型处理一维时间序列数据,参考DeepSeek给的流程编写了自己需要的代码。

关键注意事项(deepseek)

模型兼容性:

  • 避免使用RKNN不支持的算子(如自定义CUDA算子)
  • 确保 ONNX 版本与 RKNN Toolkit 兼容
  • 建议使用较新的 ONNX 版本
  • 若opset_version过高有问题,使用torch.onnx.export时可设置opset_version=10~12

量化校准:

  • 量化可显著提升速度但可能降低精度
  • 校准数据集需包含50~100张典型场景图片

开发板设置:

  • 使用散热片防止NPU过热降频
  • 通过core_mask参数分配NPU核心:
    1
    2
    
      # 使用多核提升性能
      rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0_1_2)
    

性能调优:

  • 启用target_platform=”rk3588s”适配不同变体

可能遇到的问题

error,ONNX使用了opset 19,而当前环境中的ONNX Runtime版本只支持到 opset 17

1
2
3
4
5
6
7
8
9
10
11
onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument:
[ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Failed to load model with error:
/onnxruntime_src/onnxruntime/core/graph/model_load_utils.h:57
void onnxruntime::model_load_utils::ValidateOpsetForDomain(const std::
unordered_map<std::__cxx11::basic_string<char>, int>&, const onnxruntime::
logging::Logger&, bool, const string&, int) ONNX Runtime only *guarantees*
support for models stamped with official released onnx opset versions.
Opset 19 is under development and support for this is limited.
The operator schemas and or other functionality may change before
next ONNX release and in this case ONNX Runtime will not guarantee backward
compatibility. Current official support for domain ai.onnx is till opset 17.

但我实际使用的是opset 15,opset是通过torch.onnx.export(opset_version=15)设置。

解决方法

  1. 显式指定所有算子版本(推荐)
    在导出 ONNX 时,强制所有算子使用 opset 15:
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    
    # 修改您的导出代码
    torch.onnx.export(
     model,
     dummy_input,
     "model.onnx",
     opset_version=15,
     # 添加以下参数强制所有算子使用opset 15
     custom_opsets={"": 15, "ai.onnx": 15, "ai.onnx.ml": 15},
     input_names=["input"],
     output_names=["output"],
     dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}}
    )
    
  2. 检查并降级 ONNX 模型
    使用以下代码检查实际 opset 版本:
    1
    2
    3
    4
    5
    
    import onnx
    model = onnx.load("model.onnx")
    print("Model opset versions:")
    for imp in model.opset_import:
     print(f"  Domain: {imp.domain}, Version: {imp.version}")
    

    如果发现有高于 15 的版本,可以使用 ONNX 版本转换器降级:

    1
    2
    3
    4
    5
    6
    7
    8
    
    import onnx
    from onnx import version_converter
    # 加载原始模型
    model = onnx.load("model.onnx")
    # 转换为 opset 15
    converted_model = version_converter.convert_version(model, 15)
    # 保存转换后的模型
    onnx.save(converted_model, "model_opset15.onnx")
    
  3. 更新 ONNX Runtime 和 RKNN Toolkit(我使用过)
    1
    2
    3
    4
    
    # 升级 ONNX Runtime
    pip install --upgrade onnxruntime
    # 升级 RKNN Toolkit(检查最新版本)
    pip install --upgrade rknn-toolkit2
    

    使用过:升级过ONNX Runtime。

  4. 使用 ONNX Simplifier 优化模型
    有时简化模型可以解决版本问题:
    1
    2
    
    pip install onnx-simplifier
    python -m onnxsim model.onnx model_simplified.onnx
    
本文由作者按照 CC BY 4.0 进行授权