mirror of https://github.com/coqui-ai/TTS.git
parent
e78911bf88
commit
10524b885c
25
train.py
25
train.py
|
@ -400,17 +400,20 @@ def main(args):
|
||||||
|
|
||||||
if args.restore_path:
|
if args.restore_path:
|
||||||
checkpoint = torch.load(args.restore_path)
|
checkpoint = torch.load(args.restore_path)
|
||||||
model.load_state_dict(checkpoint['model'])
|
try:
|
||||||
# Partial initialization: if there is a mismatch with new and old layer, it is skipped.
|
model.load_state_dict(checkpoint['model'])
|
||||||
# 1. filter out unnecessary keys
|
except:
|
||||||
pretrained_dict = {
|
model_dict = model.state_dict()
|
||||||
k: v
|
# Partial initialization: if there is a mismatch with new and old layer, it is skipped.
|
||||||
for k, v in checkpoint['model'].items() if k in model_dict
|
# 1. filter out unnecessary keys
|
||||||
}
|
pretrained_dict = {
|
||||||
# 2. overwrite entries in the existing state dict
|
k: v
|
||||||
model_dict.update(pretrained_dict)
|
for k, v in checkpoint['model'].items() if k in model_dict
|
||||||
# 3. load the new state dict
|
}
|
||||||
model.load_state_dict(model_dict)
|
# 2. overwrite entries in the existing state dict
|
||||||
|
model_dict.update(pretrained_dict)
|
||||||
|
# 3. load the new state dict
|
||||||
|
model.load_state_dict(model_dict)
|
||||||
if use_cuda:
|
if use_cuda:
|
||||||
model = model.cuda()
|
model = model.cuda()
|
||||||
criterion.cuda()
|
criterion.cuda()
|
||||||
|
|
Loading…
Reference in New Issue