.onnx模型转.rknn模型方法详解

首先安装对应的onnx和rknn相关环境

将所需环境安装好以后就可以新建一个:onnx2rknn.py的文件,用于模型转换。代码:

from rknn.api import RKNN
import os

if __name__ == '__main__':
    platform = 'rk3588'

    '''step 1: create RKNN object'''
    rknn = RKNN()

    '''step 2: load the .onnx model'''
    rknn.config(target_platform='rk3588')
    print('--> Loading model')
    ret = rknn.load_onnx('your onnx model.onnx')
    if ret != 0:
        print('load model failed')
        exit(ret)
    print('done')

    '''step 3: building model'''
    print('-->Building model')
    ret = rknn.build(do_quantization=False)
    if ret != 0:
        print('build model failed')
        exit()
    print('done')

    '''step 4: export and save the .rknn model'''
    OUT_DIR = 'rknn_models'
    RKNN_MODEL_PATH = './{}/actor_simple.rknn'.format(OUT_DIR)
    if not os.path.exists(OUT_DIR):
        os.mkdir(OUT_DIR)
    print('--> Export RKNN model: {}'.format(RKNN_MODEL_PATH))
    ret = rknn.export_rknn(RKNN_MODEL_PATH)
    if ret != 0:
        print('Export rknn model failed.')
        exit(ret)
    print('done')

    '''step 5: release the model'''
    rknn.release()

将上面的your onnx model换成自己的onnx名称,运行这段代码,既可在工作目录下生成一个名为rknn_models的文件夹,生成的rknn模型就在这个文件夹下。

你可能感兴趣的:(前端,linux,javascript)