zhifu gao
2024-04-25 80bd14e6bbb7bb282ff3832194648dc4a16157ca
funasr/train_utils/model_summary.py
@@ -1,4 +1,3 @@
import humanfriendly
import numpy as np
import torch
@@ -56,15 +55,8 @@
    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