首先安装对应的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模型就在这个文件夹下。