Function _onnx_opt_for_encdec

funasr/utils/export_utils.py:306–375

Internal: onnx opt for encdec. Args: model: Model instance or model name. path: TODO. enable_fp16: TODO.

(model, path, enable_fp16)

def _onnx_opt_for_encdec(model, path, enable_fp16):

    # Get input data
    # TODO: better to use real data
    """Internal: onnx opt for encdec.
    
        Args:
            model: Model instance or model name.
            path: TODO.
            enable_fp16: TODO.
        """
    input_data = model.export_dummy_inputs()

    if isinstance(input_data, torch.Tensor):
        input_data = input_data.cuda()
    else:
        input_data = tuple([i.cuda() for i in input_data])

    # Get input data for decoder module
    decoder_inputs = list()

    def get_input_hook(m, x):
        """Get input hook.
        
            Args:
                m: TODO.
                x: TODO.
            """
        decoder_inputs.extend(list(x))

    hook = model.decoder.register_forward_pre_hook(get_input_hook)
    model = model.cuda()
    model(*input_data)
    hook.remove()

    # Prevent FP16 overflow
    if enable_fp16:
        _rescale_encoder_model(model, input_data)

    fp32_model_path = f"{path}/{model.export_name}_hook.onnx"
    print("*" * 50)
    print(f"[_onnx_opt_for_encdec(fp32)]: {fp32_model_path}\n\n")
    if not os.path.exists(fp32_model_path):

        torch.onnx.export(
            model,
            input_data,
            fp32_model_path,
            verbose=False,
            do_constant_folding=True,
            opset_version=13,
            input_names=model.export_input_names(),
            output_names=model.export_output_names(),
            dynamic_axes=model.export_dynamic_axes(),
        )

    # fp32 to fp16
    fp16_model_path = f"{path}/{model.export_name}_hook_fp16.onnx"