当前位置：首页 > news >正文

Ubuntu24.04 onnx 模型转 rknn

news 来源：原创 2025/6/12 16:14:15

前面的环境配置有点懒得写，教程也很多，可以自己找 rknn-toolkit2

gitee 地址：pingli/rknn-toolkit2

试了很多开源的代码，都没办法跑通，最后自己改了一版

微调后的 qwen2 模型适用

from rknn.api import RKNN
import osif __name__ == '__main__':platform = 'rk3588'rknn = RKNN()rknn.config(target_platform='rk3588')print('--> Loading model')ret = rknn.load_onnx('onnx_model_qwen2/model.onnx', #这里改成自己的 onnx 模型路径inputs=['input_ids', 'attention_mask', 'position_ids'], input_size_list=[[1, 128], [1, 128], [1, 128]])if ret != 0:print('load model failed')exit(ret)print('done')print('-->Building model')ret = rknn.build(do_quantization=False)if ret != 0:print('build model failed')exit()print('done')OUT_DIR = 'rknn_model_qwen2' #这里改成想要的 rknn 模型路径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')rknn.release()

运行这个文件就可以了~前提是把环境配好哦