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