Pytorch学习笔记(一)-神经网络参数的冻结与解冻

在迁移学习（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)