#!/usr/bin/env python # coding: utf-8 import os import re from typing import Optional, Set from rknn.api import RKNN from math import exp from sys import exit import argparse import onnxscript from onnxscript.rewriter import pattern import onnx.numpy_helper as onh import numpy as np import onnx import onnxruntime as ort from rknn.utils import onnx_edit os.chdir(os.path.dirname(os.path.abspath(__file__))) speech_length = 171 def _remove_file(path: str, *, keep: Optional[Set[str]] = None) -> None: if not path: return keep_paths: Set[str] = {os.path.abspath(item) for item in keep} if keep else set() normalized = os.path.abspath(path) if keep_paths and normalized in keep_paths: return if not os.path.exists(normalized): return try: os.remove(normalized) print(f'cleaned temp model: {normalized}') except OSError as err: print(f'warning: failed to remove {normalized}: {err}') def _with_suffix(path: str, suffix: str) -> str: stem, ext = os.path.splitext(path) return f"{stem}{suffix}{ext}" def _sanitize_name(name: str) -> str: return re.sub(r'[^0-9A-Za-z_]', '_', name) def _insert_div_node(model: onnx.ModelProto, tensor_name: str, divisor: float = 16.0) -> bool: graph = model.graph for node in graph.node: if node.op_type == 'Div' and tensor_name in node.output: return False producer_index = None output_index = None for idx, node in enumerate(graph.node): for out_idx, output in enumerate(node.output): if output == tensor_name: producer_index = idx output_index = out_idx producer_node = node break if producer_index is not None: break if producer_index is None: raise RuntimeError(f"Producer node for tensor {tensor_name} not found.") pre_div_output = f"{tensor_name}_pre_div" producer_node.output[output_index] = pre_div_output sanitized = _sanitize_name(tensor_name) const_output = f"{sanitized}_div_const" const_node_name = f"{sanitized}_DivConst" div_node_name = f"{sanitized}_Div" const_tensor = onnx.helper.make_tensor( name=f"{const_node_name}_value", data_type=onnx.TensorProto.FLOAT, dims=[], vals=[divisor], ) const_node = onnx.helper.make_node( 'Constant', inputs=[], outputs=[const_output], value=const_tensor, name=const_node_name, ) div_node = onnx.helper.make_node( 'Div', inputs=[pre_div_output, const_output], outputs=[tensor_name], name=div_node_name, ) graph.node.insert(producer_index + 1, const_node) graph.node.insert(producer_index + 2, div_node) return True def _scale_initializer(model: onnx.ModelProto, initializer_name: str, divisor: float = 16.0) -> bool: for idx, initializer in enumerate(model.graph.initializer): if initializer.name == initializer_name: data = onh.to_array(initializer).astype(np.float32, copy=False) scaled = data / divisor model.graph.initializer[idx].CopyFrom(onh.from_array(scaled, name=initializer_name)) return True return False def convert_encoder(model_path: str): rknn = RKNN(verbose=True) ONNX_MODEL = os.path.abspath(model_path) if not os.path.isfile(ONNX_MODEL): print(f'Model file not found: {model_path}') exit(1) if not ONNX_MODEL.lower().endswith('.onnx'): print(f'Model file must be an ONNX file: {model_path}') exit(1) RKNN_MODEL = os.path.splitext(ONNX_MODEL)[0] + ".rknn" DATASET = "dataset.txt" QUANTIZE = False original_model = ONNX_MODEL preserve_files: Set[str] = {original_model} print('--> Patching model to avoid overflow issue') base_model = onnx.load(ONNX_MODEL) modified = False for layer_idx in range(48, 49): for target in [ f'/encoders.{layer_idx}/feed_forward/activation/Relu_output_0', f'/encoders.{layer_idx}/norm2/Cast_output_0', ]: modified |= _insert_div_node(base_model, target, divisor=2.0) bias_scaled = False if modified: for layer_idx in range(48, 49): bias_scaled |= _scale_initializer(base_model, f'model.encoders.{layer_idx}.feed_forward.w_2.bias', divisor=2.0) div_model_path = _with_suffix(ONNX_MODEL, "_div") onnx.save(base_model, div_model_path) if os.path.exists(div_model_path): previous_model = ONNX_MODEL ONNX_MODEL = div_model_path _remove_file(previous_model, keep=preserve_files) if modified: if bias_scaled: print('done (created div-adjusted model and scaled bias)') else: print('done (created div-adjusted model; bias initializer not found)') else: print('done (div nodes already present)') #开局先给我来个大惊喜，rknn做第一步常量折叠的时候就会在这个子图里报错，所以要单独拿出来先跑一遍 #然后把这个子图的输出结果保存下来喂给rknn extract_model_path = os.path.join(os.getcwd(), "extract_model.onnx") onnx.utils.extract_model(ONNX_MODEL, extract_model_path, ['speech_lengths'], ['/make_pad_mask/Cast_2_output_0']) sess = ort.InferenceSession(extract_model_path, providers=['CPUExecutionProvider']) extract_result = sess.run(None, {"speech_lengths": np.array([speech_length], dtype=np.int64)})[0] _remove_file(extract_model_path) # 删掉模型最后的多余transpose, 速度从365ms提升到259ms edited_model_path = _with_suffix(ONNX_MODEL, "_edited") ret = onnx_edit(model = ONNX_MODEL, export_path = edited_model_path, # # 1, len, 25055 -> 1, 25055, 1, len # 这个是坏的, 我真服了， outputs_transform = {'encoder_out': 'a,b,c->a,c,1,b'}, # outputs_transform = {'encoder_out': 'a,b,c->a,c,b'}, ) if os.path.exists(edited_model_path): previous_model = ONNX_MODEL ONNX_MODEL = edited_model_path _remove_file(previous_model, keep=preserve_files) # 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") current_model_path = ONNX_MODEL ret = rknn.load_onnx( model=current_model_path, inputs=["speech", "/make_pad_mask/Cast_2_output_0"], input_size_list=[[1, speech_length, 560], [extract_result.shape[0], extract_result.shape[1]]], input_initial_val=[None, extract_result], # outputs=["output"] ) if ret != 0: print('Load model failed!') exit(ret) print('done') _remove_file(current_model_path, keep=preserve_files) # 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') # 精度分析(可选) # rknn.accuracy_analysis(inputs=["input_content.npy"], target="rk3588", device_id=None) # usage: python convert_rknn.py path/to/model.onnx if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("model_path", type=str, help="path to source ONNX model") args = parser.parse_args() convert_encoder(args.model_path)