Function _rescale_encoder_model

funasr/utils/export_utils.py:220–259

Internal: rescale encoder model. Args: model: Model instance or model name. input_data: TODO.

(model, input_data)

def _rescale_encoder_model(model, input_data):
    # Calculate absmax
    """Internal: rescale encoder model.
    
        Args:
            model: Model instance or model name.
            input_data: TODO.
        """
    absmax = torch.tensor(0).cuda()

    def stat_input_hook(m, x, y):
        """Stat input hook.
        
            Args:
                m: TODO.
                x: TODO.
                y: TODO.
            """
        val = x[0] if isinstance(x, tuple) else x
        absmax.copy_(torch.max(absmax, val.detach().abs().max()))

    encoders = model.encoder.model.encoders
    hooks = [m.register_forward_hook(stat_input_hook) for m in encoders]
    model = model.cuda()
    model(*input_data)
    for h in hooks:
        h.remove()

    # Rescale encoder modules
    fp16_scale = int(2 * absmax // 65536)
    print(f"rescale encoder modules with factor={fp16_scale}\n\n")
    model.encoder.model.encoders0.register_forward_pre_hook(
        functools.partial(_rescale_input_hook, scale=fp16_scale),
    )
    for name, m in model.encoder.model.named_modules():
        if name.endswith("self_attn"):
            m.register_forward_hook(functools.partial(_rescale_output_hook, scale=fp16_scale))
        if name.endswith("feed_forward.w_2"):
            state_dict = {k: v / fp16_scale for k, v in m.state_dict().items()}
            m.load_state_dict(state_dict)


def _bladedisc_opt_for_encdec(model, path, enable_fp16):