unexpected key "module.conv1_1.weight" in state_dict

最新推荐文章于 2025-01-14 17:46:52 发布

木易依

最新推荐文章于 2025-01-14 17:46:52 发布

阅读量4.2k

点赞数 11

CC 4.0 BY-SA版权

分类专栏： Error

本文链接：https://blog.youkuaiyun.com/Noonebirdyou/article/details/97280122

Error 专栏收录该内容

12 篇文章

订阅专栏

本文解决在PyTorch中加载模型时遇到的KeyError问题，详细分析了错误产生的原因，包括训练与预测阶段GPU使用的不一致，并提供了两种有效的解决方案：一是通过DataParallel模块确保GPU一致性；二是修改state_dict键名，去除不必要的'module.'前缀。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

torch加载模型时出现如下错误

异常位置

save_filename = '%s_net_%s.pth' % (epoch_label, network_label)
save_path = os.path.join(self.save_dir, save_filename)
network.load_state_dict(torch.load(save_path))

异常信息

File "/data/Muyi/Github/EnlightenGAN/models/base_model.py", line 54, in load_network
network.load_state_dict(torch.load(save_path))
File "/usr/local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 522, in load_state_dict
.format(name))
KeyError: 'unexpected key "module.conv1_1.weight" in state_dict'

异常原因

最终原因在此：训练时使用GPU，使用了torch.nn.DataParallel()，而此时预测没有使用GPU，即没有使用此模块导致上述异常
if len(gpu_ids) > 0:
    netG.cuda(device=gpu_ids[0])
    netG = torch.nn.DataParallel(netG, gpu_ids)

解决方案

1):加上torch.nn.DataParallel()模块,类似我的问题只需要使用GPU即可正常运行
model = torch.nn.DataParallel(model)
cudnn.benchmark = True
2):将原来字典中module.删除掉
network.load_state_dict({k.replace('module.',''):v for k,v in torch.load(save_path).items()})
更改原来代码如下，即可在CPU/GPU下都正常运行
if len(self.gpu_ids):
    network.load_state_dict(torch.load(save_path))
else:
    network.load_state_dict({k.replace('module.',''):v for k,v in torch.load(save_path).items()})