convert_vision_encoder.py

#!/usr/bin/env python3
# ztu_somemodelruntime_rknn2: prepare_inputs_embeds

from rknn.api import RKNN
import os
import numpy as np

def main():
    # 创建RKNN实例
    rknn = RKNN(verbose=True)

    # ONNX模型路径
    ONNX_MODEL = "prepare_inputs_embeds.onnx"
    # 输出RKNN模型路径
    RKNN_MODEL = "vision_encoder.rknn"

    # 配置参数
    print("--> Config model")
    ret = rknn.config(target_platform="rk3588",
                      dynamic_input=None)
    if ret != 0:
        print('Config model failed!')
        exit(ret)

    # 加载ONNX模型
    print("--> Loading model")
    ret = rknn.load_onnx(model=ONNX_MODEL,
                         inputs=['pixel_values'],
                         input_size_list=[[1, 1, 3, 384, 384]],
                         outputs=['/aligner/layers/layers.2/Add_output_0'])
    if ret != 0:
        print('Load model failed!')
        exit(ret)

    # 构建模型
    print("--> Building model")
    ret = rknn.build(do_quantization=False)
    if ret != 0:
        print('Build model failed!')
        exit(ret)

    # 导出RKNN模型
    print("--> Export RKNN model")
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export RKNN model failed!')
        exit(ret)

    print(f'Done! The converted RKNN model has been saved to: ' + RKNN_MODEL)
    rknn.release()

if __name__ == '__main__':
    main()