Dev gzf exp (#1654)

zhifu gao

2024-04-24 861147c7308b91068ffa02724fdf74ee623a909e

 funasr/utils/export_utils.py

@@ -1,11 +1,8 @@
import os
import torch

def export_onnx(model,
                data_in=None,
            quantize: bool = False,
            opset_version: int = 14,
            **kwargs):

def export_onnx(model, data_in=None, quantize: bool = False, opset_version: int = 14, **kwargs):
   model_scripts = model.export(**kwargs)
   export_dir = kwargs.get("output_dir", os.path.dirname(kwargs.get("init_param")))
   os.makedirs(export_dir, exist_ok=True)
@@ -14,7 +11,8 @@
      model_scripts = (model_scripts,)
   for m in model_scripts:
      m.eval()
      _onnx(m,
        _onnx(
            m,
            data_in=data_in,
            quantize=quantize,
            opset_version=opset_version,
@@ -25,12 +23,15 @@
   
   return export_dir
   
def _onnx(model,

def _onnx(
    model,
         data_in=None,
         quantize: bool = False,
         opset_version: int = 14,
         export_dir:str = None,
         **kwargs):
    **kwargs
):
   
   dummy_input = model.export_dummy_inputs()
   
@@ -46,21 +47,24 @@
      opset_version=opset_version,
      input_names=model.export_input_names(),
      output_names=model.export_output_names(),
      dynamic_axes=model.export_dynamic_axes()
        dynamic_axes=model.export_dynamic_axes(),
   )
   
   if quantize:
      from onnxruntime.quantization import QuantType, quantize_dynamic
      import onnx

      quant_model_path = model_path.replace(".onnx", "_quant.onnx")
      if not os.path.exists(quant_model_path):
         onnx_model = onnx.load(model_path)
         nodes = [n.name for n in onnx_model.graph.node]
         nodes_to_exclude = [m for m in nodes if 'output' in m or 'bias_encoder' in m or 'bias_decoder' in m]
            nodes_to_exclude = [
                m for m in nodes if "output" in m or "bias_encoder" in m or "bias_decoder" in m
            ]
         quantize_dynamic(
            model_input=model_path,
            model_output=quant_model_path,
            op_types_to_quantize=['MatMul'],
                op_types_to_quantize=["MatMul"],
            per_channel=True,
            reduce_range=False,
            weight_type=QuantType.QUInt8,