From 3de70601df378664905c665d327c4c9d20c81598 Mon Sep 17 00:00:00 2001
From: 游雁 <zhifu.gzf@alibaba-inc.com>
Date: 星期一, 20 五月 2024 15:17:51 +0800
Subject: [PATCH] ds
---
funasr/train_utils/average_nbest_models.py | 10 ++++++++--
docs/images/wechat.png | 0
2 files changed, 8 insertions(+), 2 deletions(-)
diff --git a/docs/images/wechat.png b/docs/images/wechat.png
index 0ae1890..4b603fe 100644
--- a/docs/images/wechat.png
+++ b/docs/images/wechat.png
Binary files differ
diff --git a/funasr/train_utils/average_nbest_models.py b/funasr/train_utils/average_nbest_models.py
index b10c231..8553baa 100644
--- a/funasr/train_utils/average_nbest_models.py
+++ b/funasr/train_utils/average_nbest_models.py
@@ -16,7 +16,7 @@
from functools import cmp_to_key
-def _get_checkpoint_paths(output_dir: str, last_n: int = 5):
+def _get_checkpoint_paths(output_dir: str, last_n: int = 5, use_deepspeed=False):
"""
Get the paths of the last 'last_n' checkpoints by parsing filenames
in the output directory.
@@ -29,7 +29,13 @@
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]]
+ checkpoint_paths = []
+ for key, value in sorted_items[:last_n]:
+ if not use_deepspeed:
+ ckpt = os.path.join(output_dir, key)
+ else:
+ ckpt = os.path.join(output_dir, key, "mp_rank_00_model_states.pt")
+
except:
print(f"{checkpoint} does not exist, avg the lastet checkpoint.")
# List all files in the output directory
--
Gitblit v1.9.1