Medix-AI/coqui-tts
Medix-AI
/
coqui-tts
mirror of https://github.com/coqui-ai/TTS.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Filter out shape mismatch laeyrs in restoriong model for finetunning

This commit is contained in:
Eren Golge 2019-02-16 03:19:26 +01:00
parent 6ea31e47df
commit 4d2c374cfd
1 changed files with 13 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
train.py
View File

@ -367,28 +367,37 @@ def main(args):
criterion = L1LossMasked() criterion = L1LossMasked()
criterion_st = nn.BCELoss() criterion_st = nn.BCELoss()
partial_init_flag = False
if args.restore_path: if args.restore_path:
checkpoint = torch.load(args.restore_path) checkpoint = torch.load(args.restore_path)
try: try:
model.load_state_dict(checkpoint['model']) model.load_state_dict(checkpoint['model'])
except: except:
print(" > Partial model initialization.") print(" > Partial model initialization.")
partial_init_flag = True
model_dict = model.state_dict() model_dict = model.state_dict()
# Partial initialization: if there is a mismatch with new and old layer, it is skipped. # Partial initialization: if there is a mismatch with new and old layer, it is skipped.
# 1. filter out unnecessary keys # 1. filter out unnecessary keys
pretrained_dict = { pretrained_dict = {
k: v k: v
for k, v in checkpoint['model'].items() if k in model_dict for k, v in checkpoint['model'].items() if k in model_dict
} }
# 2. overwrite entries in the existing state dict # 2. filter out different size layers
pretrained_dict = {
k: v
for k, v in checkpoint['model'].items() if v.numel() == model_dict[k].numel()
}
# 3. overwrite entries in the existing state dict
model_dict.update(pretrained_dict) model_dict.update(pretrained_dict)
# 3. load the new state dict # 4. load the new state dict
model.load_state_dict(model_dict) model.load_state_dict(model_dict)
print(" | > {} / {} layers are initialized".format(len(pretrained_dict), len(model_dict)))
if use_cuda: if use_cuda:
model = model.cuda() model = model.cuda()
criterion.cuda() criterion.cuda()
criterion_st.cuda() criterion_st.cuda()
optimizer.load_state_dict(checkpoint['optimizer']) if not partial_init_flag:
optimizer.load_state_dict(checkpoint['optimizer'])
for group in optimizer.param_groups: for group in optimizer.param_groups:
group['lr'] = c.lr group['lr'] = c.lr
print( print(