Pytorch加载指定层的参数

Posted 2022-04-08 木盏

Pytorch加载模型的时候，有时候并不需要加载全部的模型参数。比如，做迁移学习的时候，我们将模型A的参数迁移到模型B上，这两个模型整体结构是不一样的，只是有些layer的是一样的，直接用model_b来load模型A的参数肯定会报错。我们只是想把模型A的部分参数加载到模型B上。
Pytorch用state_dict来结构化模型的参数，所以咱们灵活操作state_dict，就可以达到灵活加载参数的目的了。关于state_dict有疑问的，可戳《Pytorch加载模型》。
一个更生动的例子就是把ImageNet的预训练模型加载到自己的少数分类上，预训练模型有1000分类，而自己的模型只有20类，其实backbone的模型是可以加载预训练模型的。
那我们为这个栗子来撰写代码：

import torch
import torchvision.models as models

model_A = models.vgg16(pretrained=True)
torch.save(model.state_dict(), 'A_weights.pth')

model_B = models.vgg16(pretrained=False, num_classes=20)

sd = torch.load('A_weights.pth')

print(sd.keys())
part_sd = k: v for k, v in sd.items() if k not in ['classifier.6.weight', 'classifier.6.bias']
model_B.state_dict().update(part_sd)

首先，我们来分析，1000分类和20分类的模型主要区别为最后一层全连接层的神经元数量不一样，所以我们只需要把最后一层的w和b单拿出来，加载前面的层即可。
首先你得知道最后一层的w和b分别叫什么名字，可以print(sd.keys())看一下，一目了然。
如果还有其他情况，咱们要找的重点就是你想加载的层，或者不想加载的层叫什么名字就好了，然后通过dict自带的update更新一下就ok了~
有疑问可留言交流~