onnx/lx/rknn_convert.py

#!/usr/bin/env python
# coding: utf-8

# In[1]:


import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
from rknn.api import RKNN
from math import exp
from sys import exit

os.chdir(os.path.dirname(os.path.abspath(__file__)))

model_name_base = "lx"

# set input length
input_len = 1024


sample_rate = 44100
print(f"当前模型输出长度: {input_len * 512 / sample_rate * 1000} ms")


def convert_flow():
    rknn = RKNN(verbose=True)

    ONNX_MODEL=f"{model_name_base}_flow.onnx"
    RKNN_MODEL=ONNX_MODEL.replace(".onnx",".rknn")
    DATASET="dataset.txt"
    QUANTIZE=False
    detailed_performance_log = True

    # pre-process config
    print('--> Config model')
    rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3)
    print('done')

    # Load ONNX model
    print('--> Loading model')
    ret = rknn.load_onnx(model=ONNX_MODEL,
                        inputs=["z_p", "y_mask", "g"],
                        input_size_list=[[1, 192, input_len], [1, 1, input_len], [1, 256, 1]])
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    print('done')

    # Build model
    print('--> Building model')
    ret = rknn.build(do_quantization=QUANTIZE, dataset=DATASET, rknn_batch_size=None)
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')

    #export
    print('--> Export RKNN model')
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export RKNN model failed!')
        exit(ret)
    print('done')

def convert_dec():
    rknn = RKNN(verbose=True)

    ONNX_MODEL=f"{model_name_base}_dec.onnx"
    RKNN_MODEL=ONNX_MODEL.replace(".onnx",".rknn")
    DATASET="dataset.txt"
    QUANTIZE=False
    detailed_performance_log = True

    # pre-process config
    print('--> Config model')
    rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3)
    print('done')

    # Load ONNX model
    print('--> Loading model')
    ret = rknn.load_onnx(model=ONNX_MODEL,
                        inputs=["z_in", "g"],
                        input_size_list=[[1, 192, input_len], [1, 256, 1]])
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    print('done')

    # Build model
    print('--> Building model')
    ret = rknn.build(do_quantization=QUANTIZE, dataset=DATASET, rknn_batch_size=None)
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')

    #export
    print('--> Export RKNN model')
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export RKNN model failed!')
        exit(ret)
    print('done')

convert_flow()
convert_dec()