mirrors/icefall
1
0
Fork 0
mirror of https://github.com/k2-fsa/icefall.git synced 2025-08-09 01:52:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix checkpoint-writing

Browse Source
This commit is contained in:
Daniel Povey 2022-04-05 11:19:40 +08:00
parent 47d49f29d7
commit 1548cc7462
1 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
egs/librispeech/ASR/pruned_transducer_stateless2/train.py
View File

@ -376,6 +376,7 @@ def load_checkpoint_if_available(
params: AttributeDict, params: AttributeDict,
model: nn.Module, model: nn.Module,
optimizer: Optional[torch.optim.Optimizer] = None, optimizer: Optional[torch.optim.Optimizer] = None,
scheduler: Optional[torch.optim.lr_scheduler._LRScheduler] = None,
) -> Optional[Dict[str, Any]]: ) -> Optional[Dict[str, Any]]:
"""Load checkpoint from file. """Load checkpoint from file.
@ -395,6 +396,8 @@ def load_checkpoint_if_available(
The training model. The training model.
optimizer: optimizer:
The optimizer that we are using. The optimizer that we are using.
scheduler:
The scheduler that we are using.
Returns: Returns:
Return a dict containing previously saved training info. Return a dict containing previously saved training info.
""" """
@ -411,6 +414,7 @@ def load_checkpoint_if_available(
filename, filename,
model=model, model=model,
optimizer=optimizer, optimizer=optimizer,
scheduler=scheduler,
) )
keys = [ keys = [
@ -784,6 +788,7 @@ def run(rank, world_size, args):
gamma=0.5) gamma=0.5)
if checkpoints and "optimizer" in checkpoints: if checkpoints and "optimizer" in checkpoints:
logging.info("Loading optimizer state dict") logging.info("Loading optimizer state dict")
optimizer.load_state_dict(checkpoints["optimizer"]) optimizer.load_state_dict(checkpoints["optimizer"])
@ -792,7 +797,6 @@ def run(rank, world_size, args):
logging.info("Loading scheduler state dict") logging.info("Loading scheduler state dict")
scheduler.load_state_dict(checkpoints["scheduler"]) scheduler.load_state_dict(checkpoints["scheduler"])
if params.print_diagnostics: if params.print_diagnostics:
opts = diagnostics.TensorDiagnosticOptions( opts = diagnostics.TensorDiagnosticOptions(
2 ** 22 2 ** 22
@ -881,6 +885,7 @@ def run(rank, world_size, args):
params=params, params=params,
model=model, model=model,
optimizer=optimizer, optimizer=optimizer,
scheduler=scheduler,
sampler=train_dl.sampler, sampler=train_dl.sampler,
rank=rank, rank=rank,
) )