pytorch加载模型时报错

[转自]https://www.lizenghai.com/archives/4610.html

pytorch加载模型时报错

RuntimeError: Error(s) in loading state_dict for DataParallel:

Missing key(s) in state_dict: “module.backbone.layers.0.stage_1.layers.0.weight”,

这是因为加载的预训练模型之前使用了torch.nn.DataParallel()，而此时没有使用，所以可以加上该模块或者去掉。

加上torch.nn.DataParallel()模块

model = torch.nn.DataParallel(model)

cudnn.benchmark = True

创建一个没有module.的新字典，即将原来字典中module.删除掉

model.load_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():
	 # remove `module.
	 name = k[7:] 
     new_state_dict[name] = v
model.load_state_dict(torch.load(new_state_dict))

model.load_state_dict({k.replace(‘module.’,”):v for k,v in torch.load(‘model/path’).items()})