Pytorch model.load_state_dict()

1. 问题描述：

当运行到model.load_state_dict(torch.load(trained_model.pth)["state_dict"])时报错：
Missing key(s) in state_dict: "conv1.weight", "bn1.weight", "bn1.bias", "bn1.running_mean", "bn1.running_var"...
Unexpected key(s) in state_dict: "module.conv1.weight", "module.bn1.weight", "module.bn1.bias", "module.bn1.running_mean", "module.bn1.running_var"...
规律是： 缺少的key比多出的keys少一个module. 其余的一样

 

2. 原因：
训练的时候是用多个GPU训，调用了函数
model = torch.nn.DataParallel(model, device_ids=cfg.GPUS).cuda()
这样训练出来的模型存储时key会加上model.
但是我在test.py文件中并有调用函数torch.nn.DataParallel()，导致keys不一致

 

3.  解决办法：
在test.py中定义完函数后同样调用函数
model = torch.nn.DataParallel(model, device_ids=cfg.GPUS).cuda()

 

参考：https://discuss.pytorch.org/t/solved-keyerror-unexpected-key-module-encoder-embedding-weight-in-state-dict/1686