From 7e9696f15635a83ff810cc86b79c0932cb4c8a27 Mon Sep 17 00:00:00 2001
From: Djraemon <154407769+Djraemon@users.noreply.github.com>
Date: 星期一, 04 十一月 2024 11:04:52 +0800
Subject: [PATCH] Fix audio format 2.0 (#2186)
---
funasr/train_utils/model_summary.py | 18 +++++++++++++-----
1 files changed, 13 insertions(+), 5 deletions(-)
diff --git a/funasr/train_utils/model_summary.py b/funasr/train_utils/model_summary.py
index 1001160..842cd21 100644
--- a/funasr/train_utils/model_summary.py
+++ b/funasr/train_utils/model_summary.py
@@ -47,17 +47,25 @@
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"
- )
+ message += f" Number of trainable parameters: {num_params} ({percent_trainable}%)\n"
dtype = next(iter(model.parameters())).dtype
message += f" Type: {dtype}"
--
Gitblit v1.9.1