Function _bladedisc_opt_for_encdec

funasr/utils/export_utils.py:262–303

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

(model, path, enable_fp16)

def _bladedisc_opt_for_encdec(model, path, enable_fp16):
    # Get input data
    # TODO: better to use real data
    """Internal: bladedisc 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)

    # Export and optimize encoder/decoder modules
    model.encoder = _bladedisc_opt(model.encoder, input_data[:2])
    model.decoder = _bladedisc_opt(model.decoder, tuple(decoder_inputs))
    model_script = torch.jit.trace(model, input_data)
    model_script.save(os.path.join(path, f"{model.export_name}_blade.torchscript"))


def _onnx_opt_for_encdec(model, path, enable_fp16):