Fix NameError: name 'model_dict' is not defined

Closes #91
This commit is contained in:
Yves-Noel Weweler 2018-12-31 13:29:39 +01:00
parent e78911bf88
commit 10524b885c
No known key found for this signature in database
GPG Key ID: 44C0C9435C58B88D
1 changed files with 16 additions and 13 deletions

View File

@ -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()