pytoch使用nn.DataParallel导致Missing key(s) in state_dict错误

1.nn.DataParallel

nn.DataParallel是pytorch使用多gpu训练时所使用的方法，但是使用nn.DataParallel之后，模型的读取就会有所不同。最常见的情况就是使用的预训练模型并不是在多张gpu上训练得来的，没有使用nn.DataParallel包装，但现在想要使用这个预训练模型就不能直接用下面的代码

model.load_state_dict(torch.load(save_path))

否则就会报以下错误

Missing key(s) in state_dict: "module.features.0.weight", "module.features.0.bias", ....
Unexpected key(s) in state_dict: "features.0.weight", "features.0.bias", ...

2.错误原因

使用nn.DataParallel包装后的模型参数的关键字会比没用nn.DataParallel包装的模型参数的关键字前面多一个“module.”，可以看看前面的错误，missing key 和 unexpected key的差别就在那个“module.”上。还有一种情况就是missing key 是module.features.0.weight，但unexpected key是 features.module.0.weight，就是module和feature的位置反过来了。

3.解决方法

（1）在使用nn.DataParallel之前就先读取模型，然后再使用nn.DataParallel，代码的顺序是

model.load_state_dict(torch.load(save_path))
model = nn.DataParallel(model, device_ids=[0, 1]) 

（2）如果想先用nn.DataParallel，再读取模型，要使用module.load_state_dict()命令读取，代码顺序是

model= nn.DataParallel(model, device_ids=[0, 1]) 
model.module.load_state_dict(torch.load(save_path))

（3）先使用nn.DataParallel的话，也可以手动为模型参数关键字添加“module.”或是将module和features的位置调换过来

model= nn.DataParallel(model, device_ids=[0, 1]) 

from collections import OrderedDict
new_state_dict = OrderedDict()
state_dict =savepath #预训练模型路径
for k, v in state_dict.items():
	# 手动添加“module.”
    if 'module' not in k:
        k = 'module.'+k
    else:
    # 调换module和features的位置
        k = k.replace('features.module.', 'module.features.')
    new_state_dict[k]=v

model.load_state_dict(new_state_dict)

请分辨清楚三种方法的nn.DataParallel的位置

参考

https://github.com/bearpaw/pytorch-classification/issues/27
https://github.com/kuangliu/torchcv/issues/28