fix bug

游雁

2024-06-09 b75d1e89bb2f513a79bb07e9100ba1cd2bbcf40c

 funasr/train_utils/model_summary.py

@@ -1,4 +1,3 @@
import humanfriendly
import numpy as np
import torch

@@ -48,23 +47,26 @@
def model_summary(model: torch.nn.Module) -> str:
    message = "Model structure:\n"
    message += str(model)
    tot_params = sum(p.numel() for p in model.parameters())
    num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

    tot_params, num_params = 0, 0
    for name, param in model.named_parameters():
        print(
            "name: {}, dtype: {}, device: {}, trainable: {}, shape: {}, numel: {}".format(
                name, param.dtype, param.device, param.requires_grad, param.shape, param.numel()
            )
        )
        tot_params += param.numel()
        if param.requires_grad:
            num_params += param.numel()

    percent_trainable = "{:.1f}".format(num_params * 100.0 / tot_params)
    tot_params = get_human_readable_count(tot_params)
    num_params = get_human_readable_count(num_params)
    message += "\n\nModel summary:\n"
    message += f"    Class Name: {model.__class__.__name__}\n"
    message += f"    Total Number of model parameters: {tot_params}\n"
    message += (
        f"    Number of trainable parameters: {num_params} ({percent_trainable}%)\n"
    )
    num_bytes = humanfriendly.format_size(
        sum(
            p.numel() * to_bytes(p.dtype) for p in model.parameters() if p.requires_grad
        )
    )
    message += f"    Size: {num_bytes}\n"
    message += f"    Number of trainable parameters: {num_params} ({percent_trainable}%)\n"

    dtype = next(iter(model.parameters())).dtype
    message += f"    Type: {dtype}"
    return message