114 lines
3.9 KiB
Python
Executable File
114 lines
3.9 KiB
Python
Executable File
# This module is from [WeNet](https://github.com/wenet-e2e/wenet).
|
|
|
|
# ## Citations
|
|
|
|
# ```bibtex
|
|
# @inproceedings{yao2021wenet,
|
|
# title={WeNet: Production oriented Streaming and Non-streaming End-to-End Speech Recognition Toolkit},
|
|
# author={Yao, Zhuoyuan and Wu, Di and Wang, Xiong and Zhang, Binbin and Yu, Fan and Yang, Chao and Peng, Zhendong and Chen, Xiaoyu and Xie, Lei and Lei, Xin},
|
|
# booktitle={Proc. Interspeech},
|
|
# year={2021},
|
|
# address={Brno, Czech Republic },
|
|
# organization={IEEE}
|
|
# }
|
|
|
|
# @article{zhang2022wenet,
|
|
# title={WeNet 2.0: More Productive End-to-End Speech Recognition Toolkit},
|
|
# author={Zhang, Binbin and Wu, Di and Peng, Zhendong and Song, Xingchen and Yao, Zhuoyuan and Lv, Hang and Xie, Lei and Yang, Chao and Pan, Fuping and Niu, Jianwei},
|
|
# journal={arXiv preprint arXiv:2203.15455},
|
|
# year={2022}
|
|
# }
|
|
#
|
|
|
|
import logging
|
|
import os
|
|
import re
|
|
|
|
import yaml
|
|
import torch
|
|
from collections import OrderedDict
|
|
|
|
import datetime
|
|
|
|
|
|
def load_checkpoint(model: torch.nn.Module, path: str) -> dict:
|
|
if torch.cuda.is_available():
|
|
logging.info("Checkpoint: loading from checkpoint %s for GPU" % path)
|
|
checkpoint = torch.load(path)
|
|
else:
|
|
logging.info("Checkpoint: loading from checkpoint %s for CPU" % path)
|
|
checkpoint = torch.load(path, map_location="cpu")
|
|
model.load_state_dict(checkpoint, strict=False)
|
|
info_path = re.sub(".pt$", ".yaml", path)
|
|
configs = {}
|
|
if os.path.exists(info_path):
|
|
with open(info_path, "r") as fin:
|
|
configs = yaml.load(fin, Loader=yaml.FullLoader)
|
|
return configs
|
|
|
|
|
|
def save_checkpoint(model: torch.nn.Module, path: str, infos=None):
|
|
"""
|
|
Args:
|
|
infos (dict or None): any info you want to save.
|
|
"""
|
|
logging.info("Checkpoint: save to checkpoint %s" % path)
|
|
if isinstance(model, torch.nn.DataParallel):
|
|
state_dict = model.module.state_dict()
|
|
elif isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
|
state_dict = model.module.state_dict()
|
|
else:
|
|
state_dict = model.state_dict()
|
|
torch.save(state_dict, path)
|
|
info_path = re.sub(".pt$", ".yaml", path)
|
|
if infos is None:
|
|
infos = {}
|
|
infos["save_time"] = datetime.datetime.now().strftime("%d/%m/%Y %H:%M:%S")
|
|
with open(info_path, "w") as fout:
|
|
data = yaml.dump(infos)
|
|
fout.write(data)
|
|
|
|
|
|
def filter_modules(model_state_dict, modules):
|
|
new_mods = []
|
|
incorrect_mods = []
|
|
mods_model = model_state_dict.keys()
|
|
for mod in modules:
|
|
if any(key.startswith(mod) for key in mods_model):
|
|
new_mods += [mod]
|
|
else:
|
|
incorrect_mods += [mod]
|
|
if incorrect_mods:
|
|
logging.warning(
|
|
"module(s) %s don't match or (partially match) "
|
|
"available modules in model.",
|
|
incorrect_mods,
|
|
)
|
|
logging.warning("for information, the existing modules in model are:")
|
|
logging.warning("%s", mods_model)
|
|
|
|
return new_mods
|
|
|
|
|
|
def load_trained_modules(model: torch.nn.Module, args: None):
|
|
# Load encoder modules with pre-trained model(s).
|
|
enc_model_path = args.enc_init
|
|
enc_modules = args.enc_init_mods
|
|
main_state_dict = model.state_dict()
|
|
logging.warning("model(s) found for pre-initialization")
|
|
if os.path.isfile(enc_model_path):
|
|
logging.info("Checkpoint: loading from checkpoint %s for CPU" % enc_model_path)
|
|
model_state_dict = torch.load(enc_model_path, map_location="cpu")
|
|
modules = filter_modules(model_state_dict, enc_modules)
|
|
partial_state_dict = OrderedDict()
|
|
for key, value in model_state_dict.items():
|
|
if any(key.startswith(m) for m in modules):
|
|
partial_state_dict[key] = value
|
|
main_state_dict.update(partial_state_dict)
|
|
else:
|
|
logging.warning("model was not found : %s", enc_model_path)
|
|
|
|
model.load_state_dict(main_state_dict)
|
|
configs = {}
|
|
return configs
|