pytorch加载nn.DataParallel训练的模型出现的问题

1. nn.DataParallel分布式训练后，如果直接使用torch.save(model.state_dict(), model_out_path)保存模型，等到再加载模型，可以将需要加载模型的网络也弄成分布式训练。也就是说，你训练网络时，net = torch.nn.DataParallel(net)，测试模型时也用net = torch.nn.DataParallel(net)，再加载。

2. 分布式训练时，最好使用torch.save(model.module.state_dict(), model_out_path)保存模型，这样等到需要测试网络时，加载模型时用model.load_state_dict(torch.load(PATH, map_location=device))直接加载模型。

如果1，2你都不满足。又没有弄成分布式，保存的模型时又时直接使用torch.save(model.state_dict(), model_out_path)，再去加载模型时，会报以下错误：

RuntimeError: Error(s) in loading state_dict for net
Missing key(s) in state_dict: “features_1.0.weight”,…
Unexpected key(s) in state_dict: “features_1.module.0.weight”,…

可以看到错误信息提示key中多了‘.module’，那么，只要把‘.module’移除即可：

# original saved file with DataParallel
state_dict = torch.load(model_path)
# 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.replace('.module.','.') # remove `module.`
    new_state_dict[name] = v
# load params
net.load_state_dict(new_state_dict)

参考链接：
https://www.jianshu.com/p/e96a013ab5fd
https://blog.csdn.net/CV_YOU/article/details/86670188