【KeyError: ‘unexpected key “module.*****“】【KeyError: ‘unexpected key “net.“】

解决PyTorch预训练模型加载key错误问题,

最新推荐文章于 2024-01-17 15:49:28 发布

原创最新推荐文章于 2024-01-17 15:49:28 发布 · 302 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #深度学习 #python

当尝试加载使用nn.DataParallel训练的PyTorch预训练模型时，可能遇到key前带有module.的问题。文章提供了解决方案，即通过读取模型字典，移除module.前缀并保存新的参数字典，确保正确加载模型。

下载的预训练模型显示key不对，
可能是原模型使用pytorch的并行训练nn.DataParallel，存储时套了一层net，还喜欢在key前面加module，可以替换掉：

import torch
param_path = 'pretrained/****'
new_dict =  {}
pretrained_dict = torch.load(param_path, map_location=torch.device('cpu'))
pretrained_dict = pretrained_dict['net']
for k, v in pretrained_dict.items():
    k = k.replace("module.", '') 
    new_dict[k] = v
torch.save(new_dict, params_path)