onnx转rknn3588

1.安装依rknn_toolkit

pip install rknn_toolkit_your_version

2.onnx–>rknn

2.1 命令行参数

# -*- coding:utf-8 -*-
import argparse
def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--onnx_model",
        required=True, 
        help="onnx模型文件")
    parser.add_argument(
        "--platform",
        required=True,
        help="边缘计算开发板型号")
    parser.add_argument(
        "--rknn_model",
        required=True,
        help="rknn模型文件")
    parser.add_argument(
        "--test_image_dir",
        required=True,
        help="测试图片目录")
    parser.add_argument(
        "--do_quant",
        required=False,
        default=True,
        help="设置quant")
    return parser.parse_args()

2.2 转换程序

# -*- coding:utf-8 -*-
import os
import sys
import glob
from rknn.api import RKNN
from setting import parse_arguments
# 准备测试图像
def prepareDataSet(test_image_dir):
    txt_name  = './temp.txt'
    image_list = glob.glob(os.path.join(test_image_dir, '*.png'))
    with open(txt_name, 'w') as fid:
        images = 'n'.join(image_list)
        fid.write(images)
    return txt_name
def run():
    args = parse_arguments()
    onnx_model = args.onnx_model
    platform = args.platform
    rknn_model = args.rknn_model
    do_quant = args.do_quant
    test_image_dir = args.test_image_dir
    test_image_path = prepareDataSet(test_image_dir)
    rknn = RKNN(verbose=False)
    print('--> Config model')
    rknn.config(mean_values=[[0, 0, 0]], std_values=[
                    [255, 255, 255]], target_platform=platform)
    print('done')
    print('--> Loading model')
    ret = rknn.load_onnx(model=onnx_model)
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    print('done')
    # Build model
    print('--> Building model')
    ret = rknn.build(do_quantization=do_quant, dataset=test_image_path)
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')
    # Export rknn model
    print('--> Export rknn model')
    ret = rknn.export_rknn(rknn_model)
    if ret != 0:
        print('Export rknn model failed!')
        exit(ret)
    print('done')
    # Release
    rknn.release()
    os.remove(test_image_path)
if __name__ == '__main__':
    run()