欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

unexpected key "module.encoder.embedding.weight" in state_dict解决方案

程序员文章站 2022-05-27 10:33:17
...
  • 问题:unexpected key “module.encoder.embedding.weight” in state_dict
  • 原因:保存模型用nn.DataParallel,使用该方法保存模型会用module,但你加载模型时未用nn.DataParallel
  • 解决方案:
    方法1:用nn.DataParallel方法载入模型
    方法2:创建一个新的不包含module前缀的ordered dict载入模型

例子:

# original saved file with DataParallel
state_dict = torch.load('myfile.pth')
# create new OrderedDict that does not contain 'module'.
from collections import OrderedDict
new_state_dict = OrderDict()
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)
相关标签: PyTorch