pytorch载入预训练模型报错 只载入部分层

1.打印参数名称

pytorch如何打印参数名及具体值 - 简书 (jianshu.com)

2. 加载预训练模型报错

RuntimeError: Error(s) in loading state_dict for ResNet:
        size mismatch for fc.weight: copying a param with shape torch.Size([1000, 512]) from checkpoint, the shape in current model is torch.Size([128, 512]).
        size mismatch for fc.bias: copying a param with shape torch.Size([1000]) from checkpoint, the shape in current model is torch.Size([128]).

3.解决办法

(151条消息) Pytorch遇到权重不匹配的问题_hjxu2016的博客-CSDN博客 

 

 

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)