RKNN模型转换教程示例

1. 转换为 RKNN 模型

本文档说明 .rknn 扩展名模型的评估和运行方式。TensorFlow、TensorFlow Lite、Caffe、Darknet、ONNX、PyTorch 等常见训练模型，可以使用 RKNN-Toolkit2 转换为 RKNN 模型。使用其他框架训练得到的模型，也可以先转换为 ONNX 格式，再转换为 RKNN 模型。

转换流程包括：创建 RKNN 对象、初始化 SDK 环境、通过各类 load_* 接口加载模型、通过 build 构建 RKNN 模型、通过 export_rknn 导出模型，最后通过 release 释放资源。

2. 下载模型转换 Demo

下载 YOLOv5 模型转换 Demo 以及量化用数据集，并在 Ubuntu 环境中解压。Demo 中包含带有转换脚本的 yolov5_model_convert，以及包含量化图像的 quant_dataset。

yolov5_model_convert 下载链接：

https://dl.dragonwake.com/download/rv1126b/AI/demo/model_convert/yolov5_model_convert.tar.bz2

quant_dataset 下载链接：

https://dl.dragonwake.com/download/rv1126b/AI/demo/model_convert/quant_dataset.zip

3. 进入模型转换工具 Docker 环境

执行以下命令，将工作目录挂载到 Docker 容器内的 /test。其中 /dev/bus/usb:/dev/bus/usb 用于将 USB 设备传递给容器。

docker run -t -i --privileged -v /dev/bus/usb:/dev/bus/usb -v /data/project/sales/csun/rv1126b/jp/AI/demo/model_convert:/test rknn-toolkit2:2.3.2-cp38 /bin/bash

执行成功后的状态如下图所示。

注意：/data/project/sales/csun/rv1126b/jp/AI/demo/model_convert 路径需要根据实际 Ubuntu 主机使用环境进行修改。

4. 模型转换操作说明

模型转换测试 Demo 由 yolov5_model_convert 和 quant_dataset 两部分组成。yolov5_model_convert 中保存模型转换脚本，quant_dataset 中保存量化所使用的图像数据。

解压压缩文件。

unzip quant_dataset.zip
tar -jxvf yolov5_model_convert.tar.bz2

4.1 模型转换 Demo 目录结构

在工作目录中确认转换脚本和量化数据集已经正确放置。

yolov5_model_convert 文件夹主要包含以下内容。

文件	说明
best.onnx	测试用模型
get_list.py	生成量化图像列表的脚本
rknn_convert.py	模型转换脚本

4.2 生成量化图像列表

在 Docker 环境中进入模型转换工作目录，执行 gen_list.py 生成量化图像列表。

cd /test/yolov5_model_convert
python gen_list.py

执行 gen_list.py 后，将生成量化图像列表。

生成的量化图像列表如下图所示。

4.3 将 ONNX 模型转换为 RKNN 模型

执行转换脚本，将 ONNX 格式的 YOLOv5 模型转换为 RKNN 格式。转换时需要指定量化数据集，并根据目标平台和量化设置生成 RKNN 模型。

rknn_convert.py 脚本默认执行 int8 量化。脚本代码如下。

import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
from rknn.api import RKNN

ONNX_MODEL = 'best.onnx'
RKNN_MODEL = './yolov5_mask_rk3576.rknn'
DATASET = './pic_path.txt'

QUANTIZE_ON = True


if __name__ == '__main__':

    # Create RKNN object
    rknn = RKNN(verbose=True)

    if not os.path.exists(ONNX_MODEL):
        print('model not exist')
        exit(-1)

    # pre-process config
    print('--> Config model')
    rknn.config(mean_values=[[0, 0, 0]],
                std_values=[[255, 255, 255]],
                target_platform='rk3576')
    print('done')

    # Load ONNX model
    print('--> Loading model')
    ret = rknn.load_onnx(model=ONNX_MODEL)
    if ret != 0:
        print('Load yolov5 failed!')
        exit(ret)
    print('done')

    # Build model
    print('--> Building model')
    ret = rknn.build(do_quantization=QUANTIZE_ON, dataset=DATASET)
    if ret != 0:
        print('Build yolov5 failed!')
        exit(ret)
    print('done')

    # Export RKNN model
    print('--> Export RKNN model')
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export yolov5rknn failed!')
        exit(ret)
    print('done')

将 ONNX 模型 best.onnx 放置到 yolov5_model_convert 目录，执行 rknn_convert.py 脚本进行模型转换。

python rknn_convert.py

生成的模型如下图所示。该模型可在 RKNN 环境和开发板环境中运行。