记录bug——pytorch权重参数不匹配

场景:

自己修改了模型之后遇到权重参数不匹配的问题:
终端出现问题描述如下

size mismatch for fc.weight: copying a param with shape torch.Size([1000, 2048]) from checkpoint, the shape in current model is torch.Size([2, 2048]).
size mismatch for fc.bias: copying a param with shape torch.Size([1000]) from checkpoint, the shape in current model is torch.Size([2]).

原因分析

这个是因为下载的预训练模型中的全连接层是1000类别的,而当前代码类别只有2类,所以会报不匹配的错误。

解决方案:

从报错信息可以看出,是fc层的权重参数不匹配,那我们只要不load 这一层的参数就可以了。

net = se_resnet50(num_classes=2)
pretrained_dict = torch.load("./senet/seresnet50-60a8950a85b2b.pkl")

model_dict = net.state_dict()
# 重新制作预训练的权重,主要是减去参数不匹配的层,楼主这边层名为“fc”
pretrained_dict = {
    
    k: v for k, v in pretrained_dict.items() if (k in model_dict and 'fc' not in k)}
# 更新权重
model_dict.update(pretrained_dict)
net.load_state_dict(model_dict)

猜你喜欢

转载自blog.csdn.net/baobao135/article/details/129208772
今日推荐