From 6fa8ee48e117fa9c3bef450e02776e8c26b278e7 Mon Sep 17 00:00:00 2001
From: zhifu gao <zhifu.gzf@alibaba-inc.com>
Date: 星期六, 30 三月 2024 10:13:42 +0800
Subject: [PATCH] Dev gzf new (#1567)
---
funasr/train_utils/average_nbest_models.py | 10 +++++-----
1 files changed, 5 insertions(+), 5 deletions(-)
diff --git a/funasr/train_utils/average_nbest_models.py b/funasr/train_utils/average_nbest_models.py
index c31c1cb..3413461 100644
--- a/funasr/train_utils/average_nbest_models.py
+++ b/funasr/train_utils/average_nbest_models.py
@@ -30,6 +30,7 @@
sorted_items = sorted_items[:last_n] if avg_keep_nbest_models_type == "acc" else sorted_items[-last_n:]
checkpoint_paths = [os.path.join(output_dir, key) for key, value in sorted_items[:last_n]]
except:
+ print(f"{checkpoint} does not exist, avg the lastet checkpoint.")
# List all files in the output directory
files = os.listdir(output_dir)
# Filter out checkpoint files and extract epoch numbers
@@ -56,10 +57,9 @@
state_dicts.append(torch.load(path, map_location='cpu')['state_dict'])
else:
print(f"Checkpoint file {path} not found.")
- continue
# Check if we have any state_dicts to average
- if not state_dicts:
+ if len(state_dicts) < 1:
raise RuntimeError("No checkpoints found for averaging.")
# Average or sum weights
@@ -75,6 +75,6 @@
# Perform average for other types of tensors
stacked_tensors = torch.stack(tensors)
avg_state_dict[key] = torch.mean(stacked_tensors, dim=0)
-
- torch.save({'state_dict': avg_state_dict}, os.path.join(output_dir, f"model.pt.avg{last_n}"))
- return avg_state_dict
\ No newline at end of file
+ checkpoint_outpath = os.path.join(output_dir, f"model.pt.avg{last_n}")
+ torch.save({'state_dict': avg_state_dict}, checkpoint_outpath)
+ return checkpoint_outpath
\ No newline at end of file
--
Gitblit v1.9.1