博客

1、加载指定预训练权重

有时无法直接利用官方的预训练模型，需要自己手动将指定模型进行预训练。主要有以下语句作为参考：

model = models.resnet34(pretrained=False)
pretrained_dict = torch.load('./pretrain/resnet34-333f7ec4.pth')
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model.state_dict()}
model.load_state_dict(pretrained_dict)
model.fc = torch.nn.Linear(512, 5) # 512为原始fc的数目，5是自己任务的分类数

在自己的任务，分类类别数不一致时，通常出现如下类似的错误：

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


2、移除指定层，再更新权重

只是不加载fc层也可以，与上面用法一定程度等价。

model = models.resnet34(pretrained=False)
pretrained_dict = torch.load('./pretrain/resnet34-333f7ec4.pth')
model_dict = model.state_dict()
pretrained_dict = {k: v for k, v in pretrained_dict.items() if (k in model_dict and 'fc' not in k)} # 将'fc'这一层的权重选择不加载即可。
model_dict.update(pretrained_dict) # 更新权重
model.load_state_dict(model_dict)