speech_recognition/FunASR
1
0
Fork 0
mirror of https://github.com/modelscope/FunASR synced 2025-09-15 14:48:36 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

ds

Browse Source
This commit is contained in:
游雁 2024-05-20 13:40:26 +08:00
parent b3b1015809
commit a19ee46d2c
2 changed files with 6 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
funasr/train_utils/average_nbest_models.py
View File

@ -62,7 +62,8 @@ def average_checkpoints(output_dir: str, last_n: int = 5, **kwargs):
# Check if we have any state_dicts to average # Check if we have any state_dicts to average
if len(state_dicts) < 1: if len(state_dicts) < 1:
raise RuntimeError("No checkpoints found for averaging.") print("No checkpoints found for averaging.")
return
# Average or sum weights # Average or sum weights
avg_state_dict = OrderedDict() avg_state_dict = OrderedDict()

10
funasr/train_utils/trainer_ds.py
View File

@ -168,8 +168,7 @@ class Trainer:
""" """
step_in_epoch = None if step is None else step_in_epoch step_in_epoch = None if step is None else step_in_epoch
if self.use_deepspeed: if self.use_deepspeed:
with torch.no_grad():
model.save_checkpoint(save_dir=model_dir, tag=tag, client_state=info_dict)
logging.info(f"Save checkpoint: {epoch}, rank: {self.local_rank}\n") logging.info(f"Save checkpoint: {epoch}, rank: {self.local_rank}\n")
# self.step_or_epoch += 1 # self.step_or_epoch += 1
state = { state = {
@ -273,8 +272,7 @@ class Trainer:
elif self.use_fsdp: elif self.use_fsdp:
pass pass
step_in_epoch = None if step is None else step_in_epoch elif self.rank == 0:
if self.rank == 0:
logging.info(f"Save checkpoint: {epoch}, rank: {self.local_rank}\n") logging.info(f"Save checkpoint: {epoch}, rank: {self.local_rank}\n")
# self.step_or_epoch += 1 # self.step_or_epoch += 1
state = { state = {
@ -385,7 +383,7 @@ class Trainer:
if self.use_deepspeed: if self.use_deepspeed:
ckpt = os.path.join(self.output_dir, "model.pt") ckpt = os.path.join(self.output_dir, "model.pt")
if os.path.isfile(ckpt): if os.path.exists(ckpt):
_, checkpoint = model_engine.load_checkpoint(self.output_dir, "model.pt") _, checkpoint = model_engine.load_checkpoint(self.output_dir, "model.pt")
self.saved_ckpts = checkpoint["saved_ckpts"] self.saved_ckpts = checkpoint["saved_ckpts"]
@ -712,7 +710,7 @@ class Trainer:
"data_split_num": kwargs.get("data_split_num", 1), "data_split_num": kwargs.get("data_split_num", 1),
"log_step": batch_idx + kwargs.get("start_step", 0), "log_step": batch_idx + kwargs.get("start_step", 0),
"batch_total": batch_idx, "batch_total": batch_idx,
"step_in_epoch": step_in_epoch, "step_in_epoch": batch_idx,
"lr": 0.0, "lr": 0.0,
} }