Pytorch加载模型时报错

在Pytorch中遇到加载模型时的'Unexpected key(s) in state_dict'错误,问题源于多GPU训练保存的模型带有'module.'前缀。解决方案包括在加载时使用'strict=False'参数或手动修改权重键值对,移除'module.'前缀,确保模型与权重匹配。

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

加载模型时出现Unexpected key(s) in state_dict错误

报错截图如下:
在这里插入图片描述
反复排查问题没发现为何如此,查看pytorch中文文档发现保存和加载模型方法都完全正确,模型保存和加载代码对比中文文档截图如下:
在这里插入图片描述
其中一个方法是在加载模型时添加参数strict=False,可以只保留键值相同的参数避免出错,用法如下:model.load_state_dict(ckpt[‘state_dict’],False),注意ckpt是我加载权重时自定义的对象名。
但是继续训练发现实际上原有的权重都没有加载进来,因为没有键值相同的参数。此处行不通

继续深究发现多一个module是因为通过多GPU训练保存的模型!如果是在cpu训练的模型则不会带有module,解决很简单,在加载权重前先将模型映射到GPU训练,
model = torch.nn.DataParallel(model)
ckpt = torch.load(args.weights)
model.load_state_dict(ckpt[‘state_dict’])
或者还有解决办法,直接将权重里面key值中含有的module.删掉,改写权重键值对
ckpt = torch.load(model_path,map_location=‘cpu’)
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值