[转自]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()})