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(),而此时没有使用,所以可以加上该模块或者去掉。
1,加上torch.nn.DataParallel()模块
model = torch.nn.DataParallel(model)
cudnn.benchmark = True
2、创建一个没有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
# load params
model.load_state_dict(torch.load(new_state_dict))
3、思想同2
model.load_state_dict({k.replace('module.',''):v for k,v in torch.load('model/path').items()})
更多讨论请参考:
https://discuss.pytorch.org/t/solved-keyerror-unexpected-key-module-encoder-embedding-weight-in-state-dict/1686/3
记录pytorch .pth模型转onnx遇到的问题:
1、Missing key(s) in state_dict:
解决方法:
# original saved file with DataParallel
state_dict = torch.load('myfile.pth.tar')
# 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[7:] # remove `module.`
new_state_dict[name] = v
# load params
model.load_state_dict(new_state_dict)
2、RuntimeError: Expected object of backend CPU but got backend CUDA for sequence element 2 in sequence argument at position #1 'tensors'
解决方法:
在GPU上训练的模型在CPU上运行时会出现的问题
加上
model = torch.load('model.pth',map_location='cpu')
3、AttributeError: 'xxxx' object has no attribute 'seek'. You can only torch.load from a file that is seekable. Please pre-load the data into a buffer like io.BytesIO and try to load from it instead.
网友评论