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