瑞芯微 Rockchip 系列 RK3588 主流深度学习框架模型转成 rknn 模型教程

前言

在瑞芯微 Rockchip 芯片上进行 NPU 推理，需要先将模型文件转换成 rknn 模型文件，才能执行各种推理任务。本文将介绍如何安装各种工具，并最终实现将各种深度学习框架的模型文件转换成 rknn 文件。

本教程不仅适合 RK3588 平台，也适合其他 RK 系列平台，例如 RK3566、RK3568、RK3562 等。具体平台请参考 RKNN-Toolkit2 文档。

安装 RKNN-Toolkit2 环境

RKNN-Toolkit2 是瑞芯微官方提供的工具包，用于将各种深度学习框架的模型文件转换成 rknn 文件。

安装 Conda

RKNN-Toolkit2 转换依赖于 Python 3.8（建议版本），而一般我们使用 Conda 来管理 python 环境，所以需要先安装 Conda。

下面的命令会下载 Anaconda 的安装包，并执行安装。

wget https://repo.anaconda.com/archive/Anaconda3-2024.10-1-Linux-x86_64.sh
bash Anaconda3-2024.10-1-Linux-x86_64.sh

安装完成后，运行下面命令，查看是否安装成功。

conda --version

使用 Conda 创建 Python 3.8 环境

通过以下命令创建名称为 toolkit2 的 Python 3.8 环境：

conda create -n toolkit2 python=3.8

激活 toolkit2 环境，后续将在此环境中安装 RKNN-Toolkit2：

conda activate toolkit2
# 成功后，命令行提示符会变成以下形式：
# (toolkit2) xxx@xxx:~$

安装 RKNN-Toolkit2

在终端运行如下命令：

pip install rknn-toolkit2 -i https://pypi.org/simple

运行下面命令验证是否安装成功：

# 进入 Python 交互模式
python

# 导入 RKNN 类
from rknn.api import RKNN

如果没有输出错误，则表示 RKNN-Toolkit2 安装成功。

模型转换

模型转换流程

RKNN-Toolkit2 的转换流程如下：

可以看到，整个流程非常简单，初始化 -> 配置 -> 加载 -> 构建 -> 导出模型 -> 释放资源。

目前RKNN-Toolkit2支持多个主流深度学习框架的模型转换，包括：

• Caffe（推荐版本为1.0）
• TensorFlow（推荐版本为1.12.0~2.8.0）
• TensorFlow Lite（推荐版本为Schema version = 3）
• ONNX（推荐版本为1.7.0~1.10.0）
• PyTorch（推荐版本为1.6.0~1.13.1）
• Darknet（推荐版本为Commit ID = 810d7f7）

在 RKNN 模型转换过程中，可以选择是否要进行模型量化。量化的作用是减少模型的大小，提高推理速度，但可能会牺牲一定的精度。如果要量化，需要传入量化校准的数据集，这是为了确保量化后的模型在推理时，能够达到相对较好的精度。

编写模型转换示例

下面我以转换 yolo11 的 ONNX 模型为例，手把手带大家编写一个简单的模型转换示例，请确保上面的 Conda 已经安装并且 toolkit2 激活环境。

在最下面我也会提供整个示例的完整代码仓库，大家可以直接下载使用。

首先，在 Linux 终端运行如下命令：

mkdir rknn_trans_test
cd rknn_trans_test

# 下载 yolo11 的 ONNX 模型
wget -O ./yolo11n.onnx https://ftrg.zbox.filez.com/v2/delivery/data/95f00b0fc900458ba134f8b180b3f7a1/examples/yolo11/yolo11n.onnx

接着下载量化的数据集

# 下载 yolo11 的量化数据集
wget https://github.com/qaz624824554/rknn_trans_test/raw/refs/heads/master/COCO.tar.gz
tar -xzvf COCO.tar.gz

创建 convert.py 文件

touch convert.py

将下面的代码复制到 convert.py 文件中：

import sys
from rknn.api import RKNN

DATASET_PATH = './COCO/coco_subset_20.txt'  # 量化数据集路径
DEFAULT_RKNN_PATH = './yolo11n.rknn'        # 默认输出RKNN模型路径
DEFAULT_QUANT = True                         # 默认是否量化

def parse_arg():
    """
    解析命令行参数
    Returns:
        model_path: ONNX模型路径
        platform: 目标平台
        do_quant: 是否进行量化
        output_path: RKNN模型输出路径
    """
    if len(sys.argv) < 3:
        print("Usage: python3 {} onnx_model_path [platform] [dtype(optional)] [output_rknn_path(optional)]".format(sys.argv[0]))
        print("       platform choose from [rk3562,rk3566,rk3568,rk3588,rk3576,rk1808,rv1109,rv1126]")
        print("       dtype choose from [i8, fp] for [rk3562,rk3566,rk3568,rk3588,rk3576]")
        print("       dtype choose from [u8, fp] for [rk1808,rv1109,rv1126]")
        exit(1)

    model_path = sys.argv[1]
    platform = sys.argv[2]

    # 根据输入参数确定是否量化
    do_quant = DEFAULT_QUANT
    if len(sys.argv) > 3:
        model_type = sys.argv[3]
        if model_type not in ['i8', 'u8', 'fp']:
            print("ERROR: Invalid model type: {}".format(model_type))
            exit(1)
        elif model_type in ['i8', 'u8']:
            do_quant = True
        else:
            do_quant = False

    # 确定输出路径
    if len(sys.argv) > 4:
        output_path = sys.argv[4]
    else:
        output_path = DEFAULT_RKNN_PATH

    return model_path, platform, do_quant, output_path

if __name__ == '__main__':
    # 解析命令行参数
    model_path, platform, do_quant, output_path = parse_arg()

    # 创建RKNN对象
    rknn = RKNN(verbose=False)

    # 配置预处理参数
    print('--> Config model')
    rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform=platform)
    print('done')

    # 加载ONNX模型
    print('--> Loading model')
    ret = rknn.load_onnx(model=model_path)
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    print('done')

    # 构建RKNN模型
    print('--> Building model')
    ret = rknn.build(do_quantization=do_quant, dataset=DATASET_PATH)
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')

    # 导出RKNN模型
    print('--> Export rknn model')
    ret = rknn.export_rknn(output_path)
    if ret != 0:
        print('Export rknn model failed!')
        exit(ret)
    print('done')

    # 释放资源
    rknn.release()

保存，运行如下命令

# 用法: python convert.py model_path [rk3566|rk3588|rk3562] [i8/u8/fp] [output_path]
python convert.py ./yolo11n.onnx rk3588 i8 ./yolo11n.rknn

看到下面的运行效果表示转换成功，并且可以看到当前目录下多了一个 yolo11n.rknn 文件。

关于转换过程中用到的具体函数参数，可以参考 RKNN-Toolkit2 文档，这里不再过多赘述。

示例代码仓库：https://github.com/qaz624824554/rknn_trans_test

总结

本文介绍了如何在 RK3588 平台上进行主流深度学习框架模型转换，并提供了完整的示例代码。希望对大家有所帮助。

如果大家有任何问题，欢迎在评论区留言，我会尽快回复。

如果觉得本文对你有帮助，欢迎点赞、收藏、转发。

谢谢大家！