在迁移学习(Transfer Learning)等领域,通常会将预训练模型(Pre-trained Model)主干部分的权重进行冻结,再在目标领域的数据集下对余下权重参数进行调参。上述过程涉及神经网络参数冻结与解冻。
在Pytorch中,网络参数通过梯度下降法及其衍生算法进行优化,其关键在于梯度的逐层传递。因此,通过关闭或打开参数的梯度计算即可实现参数的冻结与解冻。Pytorch中,网络参数均为张量(Tensor),涉及梯度计算的特征(Attribute)为requires_grad,默认情况下requires_grad=True,即默认情况下计算Tensor的梯度。
model.fc.weight.requires_grad = False
model.fc.bias.requires_grad = False
PyTorch提供了state_dict()和load_state_dict()两个方法用来保存和加载模型参数,前者将模型参数保存为字典形式,后者将字典形式的模型参数载入到模型当中。下面是使用预训练权重(加载预训练模型)的代码,其中model_path就是预训练权重文件的路径:
# 第一步:读取当前模型参数
model_dict = model.state_dict()
# 第二步:读取预训练模型
pretrained_dict = torch.load(model_path, map_location = device)
pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) == np.shape(v)}
# 第三步:使用预训练的模型更新当前模型参数
model_dict.update(pretrained_dict)
# 第四步:加载模型参数
model.load_state_dict(model_dict)
但是,使用load_state_dict()加载模型参数时,要求保存的模型参数键值类型和模型完全一致
一旦我们对模型结构做了些许修改,就会出现类似unexpected key module.xxx.weight问题。比如在目标检测模型中,如果修改了主干特征提取网络,只要不是直接替换为现有的其它神经网络,基本上预训练权重是不能用的,要么就自己判断权值里卷积核的shape然后去匹配,要么就只能利用这个主干网络在诸如ImageNet这样的数据集上训练一个自己的预训练模型;如果修改的是后面的neck或者是head的话,前面的backbone的预训练权重还是可以用的。下面是权值匹配的示例代码,把不匹配的直接pass了:
model_dict = model.state_dict()
pretrained_dict = torch.load(model_path, map_location=device)
temp = {}
for k, v in pretrained_dict.items():
try:
if np.shape(model_dict[k]) == np.shape(v):
temp[k]=v
except:
pass
model_dict.update(temp)