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

呆毛王的意志

于 2023-04-15 11:11:47 发布

阅读量274

点赞数

CC 4.0 BY-SA版权

文章标签： pytorch 深度学习 python

本文链接：https://blog.youkuaiyun.com/github_38575699/article/details/130167359

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

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

下载的预训练模型显示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)