深度学习 加载网络参数遇到 Missing key(s)

在PyTorch中加载预训练模型参数时遇到MissingKey错误,通常是因为网络结构定义与加载的参数不匹配。文章提到,可以通过比较模型结构和参数键来对应不同的命名。作者选择修改定义的网络结构以适应加载的参数,特别地,由于保存的参数不包含归一化层和激活函数层,以及dropout层在测试时不需要,因此作者去掉了dropout层以使参数对应。

在加载别人已经训练好的网络参数做测试时遇到以下问题
在这里插入图片描述图(1)
使用net.load_state_dict() 加载参数,遇到Missing key(s)的报错,必然是定义的网络结构与加载的网络结构不同。
如上图所示,待加载的网络参数对比定义的网络参数Missing和Unexpected一 一可以对应。从中得出,待加载的网络和定义的网络是可以对应上的,只是网络命名不一样。对于这种问题有两种方法解决,要么修改待加载的网络名,要么修改定义是的网络名。本文采用第二种。
首先打印一下定义的网络结构
print(net)
在这里插入图片描述
图(2)
再打印一下加载的网络参数关键字
state_dict=torch.load(load_path)
print(state_dict.keys())
在这里插入图片描述
图(3)
在保存网络参数时是不会保存归一化层和激活函数层的,图(3)关键字中 model.10.conv_block.1.weight 中model代表整个模型,10代表第十层,conv_block代表conv_block模块,1代表该模块的第一层,weight代表参数,与图(2)中的网络一 一对应。
再来看图(1)中的错误,图(2)中的模型第10层中的conv_block模块中第6层是卷积层,而在图(3)的网络参数中第五层是卷积层参数。所以只要将卷积层的参数、偏差一 一对应即可。
在这里我是将所有的dropout层都去掉,因为测试时dropout本身就用不到。当然也可以调整一下位置。

站内引用中未提及解决UNet2DConditionModel加载state_dict时出现Missing key(s) in state_dict错误的具体方法。不过,结合常见的深度学习模型加载问题解决思路,可尝试以下方法: 1. **检查模型定义和权重文件是否匹配**:确保使用的模型定义代码和保存权重时的模型结构一致。若模型结构有修改,可能会导致部分键缺失。 2. **手动补充缺失的键**:若缺失的键对模型影响不大,可手动为state_dict补充缺失的键,并赋予合适的初始值。示例代码如下: ```python import torch # 加载原始的state_dict original_state_dict = torch.load('your_model_weights.pth') # 检查模型定义 model = UNet2DConditionModel() # 假设这是你的模型定义 # 获取模型的所有键 model_keys = model.state_dict().keys() # 找出缺失的键 missing_keys = [key for key in model_keys if key not in original_state_dict] # 手动补充缺失的键 for key in missing_keys: # 根据键的类型和形状,赋予合适的初始值 if 'weight' in key: original_state_dict[key] = torch.randn(*model.state_dict()[key].shape) elif 'bias' in key: original_state_dict[key] = torch.zeros(*model.state_dict()[key].shape) # 加载补充后的state_dict model.load_state_dict(original_state_dict, strict=False) ``` 3. **使用`strict=False`参数**:在调用`load_state_dict`方法时,将`strict`参数设置为`False`,可忽略部分键的缺失。但要注意,这可能会使部分层未正确加载权重。示例代码如下: ```python import torch model = UNet2DConditionModel() state_dict = torch.load('your_model_weights.pth') model.load_state_dict(state_dict, strict=False) ```
评论
成就一亿技术人!
拼手气红包6.0元
还能输入1000个字符
 
红包 添加红包
表情包 插入表情
 条评论被折叠 查看
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值