Pytorch加载指定层的参数

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了~
有疑问可留言交流~

你可能感兴趣的:(Pytorch那些事儿,pytorch,load_state_dict,state_dict,Pytorch,tensorflow)