Ubuntu24.04 onnx 模型转 rknn

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

gitee 地址:pingli/rknn-toolkit2


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

微调后的 qwen2 模型适用

from rknn.api import RKNN
import os
 
if __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()

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

你可能感兴趣的:(大模型,人工智能,个人开发)