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

Use GPU for averaging checkpoints if possible.

Browse Source
This commit is contained in:
Fangjun Kuang 2021-10-18 11:32:10 +08:00
parent bd7c2f7645
commit 47519b36a4
3 changed files with 11 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
egs/librispeech/ASR/conformer_ctc/decode.py
View File

@ -638,7 +638,8 @@ def main():
if start >= 0:
filenames.append(f"{params.exp_dir}/epoch-{i}.pt")
logging.info(f"averaging {filenames}")
model.load_state_dict(average_checkpoints(filenames))
model.to(device)
model.load_state_dict(average_checkpoints(filenames, device=device))
if params.export:
logging.info(f"Export averaged model to {params.exp_dir}/pretrained.pt")

3
egs/librispeech/ASR/tdnn_lstm_ctc/decode.py
View File

@ -455,7 +455,8 @@ def main():
if start >= 0:
filenames.append(f"{params.exp_dir}/epoch-{i}.pt")
logging.info(f"averaging {filenames}")
model.load_state_dict(average_checkpoints(filenames))
model.to(device)
model.load_state_dict(average_checkpoints(filenames, device=device))
if params.export:
logging.info(f"Export averaged model to {params.exp_dir}/pretrained.pt")

10
icefall/checkpoint.py
View File

@ -120,22 +120,26 @@ def load_checkpoint(
return checkpoint
def average_checkpoints(filenames: List[Path]) -> dict:
def average_checkpoints(
filenames: List[Path], device: torch.device = torch.device("cpu")
) -> dict:
"""Average a list of checkpoints.
Args:
filenames:
Filenames of the checkpoints to be averaged. We assume all
checkpoints are saved by :func:`save_checkpoint`.
device:
Move checkpoints to this device before averaging.
Returns:
Return a dict (i.e., state_dict) which is the average of all
model state dicts contained in the checkpoints.
"""
n = len(filenames)
avg = torch.load(filenames[0], map_location="cpu")["model"]
avg = torch.load(filenames[0], map_location=device)["model"]
for i in range(1, n):
state_dict = torch.load(filenames[i], map_location="cpu")["model"]
state_dict = torch.load(filenames[i], map_location=device)["model"]
for k in avg:
avg[k] += state_dict[k]