DAMO-YOLO预训练模型转RKNN

推荐环境

ubuntu 20.04

获取rknn项目

Git 克隆rknn-toolkit2项目到本地， 连接如下： 

GitHub - rockchip-linux/rknn-toolkit2Contribute to rockchip-linux/rknn-toolkit2 development by creating an account on GitHub.https://github.com/rockchip-linux/rknn-toolkit2.git

安装依赖

conda创建一个python3.8的环境；

使用rknn环境；

进入项目doc目录，安装 requirements_cp38-1.5.2.txt 中描述的依赖库；

conda create -n rknn python=3.8 -y

conda activate rknn

cd rknn-toolkit2/doc

pip install -r requirements_cp38-1.5.2.txt

若安装TensorFlow时遇到超时的问题， 可以使用镜像源单独安装该依赖后，再安装requirements_cp38-1.5.2.txt中描述的其它依赖。

pip install tensorflow==2.8.0 -i https://pypi.douban.com/simple/

进入目录 rknn-toolkit2/packages/，使用 rknn_toolkit2-1.5.2+b642f30c-cp38-cp38-linux_x86_64.whl 安装rknn。

pip install rknn_toolkit2-1.5.2+b642f30c-cp38-cp38-linux_x86_64.whl

准备待转模型&数据集

从DAMO-YOLO 项目中直接下载已训练好的模型，或者自己训练模型。本文使用DAMO-YOLO-NI模型， 将模型存放到 rknn-toolkit2/converter/ 目录下

从 rknn-toolkit2/examples/onnx/yolov5/ 目录复制  dataset.txt 和 bus.jpg 文件至 rknn-toolkit2/pretrained/ 目录下

编写转换脚本

在 rknn-toolkit2/converter/ 目录下编写转换脚本 onnx2rknn.py 

import argparse
from rknn.api import RKNN


QUANTIZE_ON = True
DATASET = './dataset.txt'

def make_parser():
    parser = argparse.ArgumentParser()

    parser.add_argument("-t", "--target_platform", type=str, help='芯片平台名称!')
    parser.add_argument("-p", "--origin_path", type=str, help='onxx文件路径!')
    parser.add_argument("-d", "--dest_path", type=str, help='导出rknn文件目录!')

    return parser

def main():
    args = make_parser().parse_args()
    target_platform = args.target_platform
    origin_path = args.origin_path
    dest_path = args.dest_path

    print('ORIGIN-PATH:', target_platform)
    print('ORIGIN-PATH:', origin_path)
    print('DEST-PATH:', dest_path)

    # 校验输入的模型文件是否存在，并且获取输出文件名
    if os.path.isfile(origin_path):
        file_name = os.path.basename(origin_path)
        file_name = file_name[:file_name.index(".")] + ".rknn"
    else:
        print('模型文件不存在')
        exit(-1)
    print('done')

    if os.path.isdir(dest_path):
        dest_file = dest_path + "/" + file_name
    else:
        print('输出目录不存在!')
        exit(-1)
    print('done')

    print('DEST-FILE:', dest_file)

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

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

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

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

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


if __name__ == '__main__':
    main()

使用脚本转换

进入 rknn-toolkit2/converter/ 目录，运行脚本命令

python onnx2rknn.py -t rk3588 -p ./damoyolo-tinynasL20_Nl.onnx -d ./

最终得到转换后的rknn模型， damoyolo_tinynasL20_Nl.rknn

---