Missing key(s) in state_dict: Unexpected key(s) in state_dict

weixin_34104341發表於2020-04-07

如果載入的預訓練模型之前使用了torch.nn.DataParallel(),而此時的訓練並沒有使用,則會出現這樣的錯誤。
解決方案有兩個:
1:此時的訓練加入torch.nn.DataParallel()即可。
2:建立一個沒有module.的新字典,即將原來字典中module.刪除掉。
解決方案1:

model = torch.nn.DataParallel(model)
cudnn.benchmark = True


解決方案2:
# original saved file with DataParallel
state_dict = torch.load('myfile.pth')
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
    name = k[7:] # remove `module.`
    new_state_dict[name] = v
# load params
model.load_state_dict(new_state_dict)

解決方案3:
model.load_state_dict({k.replace('module.',''):v for k,v in torch.load('myfile.pth').items()})

轉載於:https://www.cnblogs.com/llfctt/p/11045066.html

相關文章