pytorch：预训练权重、冻结训练和断点恢复

知乎—吵鸡凶鸭OvO  侵删

原文 https://zhuanlan.zhihu.com/p/446812760

本篇文章将介绍神经网络训练过程中的三个必备技能：使用预训练权重、冻结训练和断点恢复，巧妙运用这三个技巧可以很有效地提高网络的训练效率和效果。

01

引言

If I have seen further, it is by standing on the shoulders of giants.

迁移学习在计算机视觉领域中是一种很流行的方法，因为它可以建立精确的模型，耗时更短。利用迁移学习，不是从零开始学习，而是从之前解决各种问题时学到的模式开始。这样，我们就可以利用以前的学习成果，避免从零开始。

02

使用预训练权重

在计算机视觉领域中，迁移学习通常是通过使用预训练模型来表示的。预训练模型是在大型基准数据集上训练的模型，用于解决相似的问题。由于训练这种模型的计算成本较高，因此，导入已发布的成果并使用相应的模型是比较常见的做法。例如，在目标检测任务中，首先要利用主干神经网络进行特征提取，这里使用的backbone一般就是VGG、ResNet等神经网络，因此在训练一个目标检测模型时，可以使用这些神经网络的预训练权重来将backbone的参数初始化，这样在一开始就能提取到比较有效的特征。

可能大家会有疑问，预训练权重是针对他们数据集训练得到的，如果是训练自己的数据集还能用吗？预训练权重对于不同的数据集是通用的，因为特征是通用的。一般来讲，从0开始训练效果会很差，因为权值太过随机，特征提取效果不明显。对于目标检测模型来说，一般不从0开始训练，至少会使用主干部分的权值，虽然有些论文提到了可以不用预训练，但这主要是因为他们的数据集比较大而且他们的调参能力很强。如果从0开始训练，网络在前几个epoch的Loss可能会非常大，并且多次训练得到的训练结果可能相差很大，因为权重初始化太过随机。

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)

03

冻结训练

冻结训练其实也是迁移学习的思想，在目标检测任务中用得十分广泛。因为目标检测模型里，主干特征提取部分所提取到的特征是通用的，把backbone冻结起来训练可以加快训练效率，也可以防止权值被破坏。在冻结阶段，模型的主干被冻结了，特征提取网络不发生改变，占用的显存较小，仅对网络进行微调。在解冻阶段，模型的主干不被冻结了，特征提取网络会发生改变，占用的显存较大，网络所有的参数都会发生改变。举个例子，如果在解冻阶段设置batch_size为4，那么在冻结阶段有可能可以把batch_size设置到8。下面是进行冻结训练的示例代码，假设前50个epoch冻结，后50个epoch解冻：

# 冻结阶段训练参数，learning_rate和batch_size可以设置大一点
Init_Epoch          = 0
Freeze_Epoch        = 50
Freeze_batch_size   = 8
Freeze_lr           = 1e-3
# 解冻阶段训练参数，learning_rate和batch_size设置小一点
UnFreeze_Epoch      = 100
Unfreeze_batch_size = 4
Unfreeze_lr         = 1e-4
# 可以加一个变量控制是否进行冻结训练
Freeze_Train        = True
# 冻结一部分进行训练
batch_size  = Freeze_batch_size
lr          = Freeze_lr
start_epoch = Init_Epoch
end_epoch   = Freeze_Epoch
if Freeze_Train:
  for param in model.backbone.parameters():
  param.requires_grad = False
# 解冻后训练
batch_size  = Unfreeze_batch_size
lr          = Unfreeze_lr
start_epoch = Freeze_Epoch
end_epoch   = UnFreeze_Epoch
if Freeze_Train:
  for param in model.backbone.parameters():
  param.requires_grad = True

如果不进行冻结训练，一定要注意参数设置，注意上述代码中冻结阶段和解冻阶段的learning_rate和batch_size是不一样的，另外起始epoch和结束epoch也要重新调整一下。如果是从0开始训练模型（不使用预训练权重），那么一定不能进行冻结训练。

04

断点恢复

在上面冻结训练和解冻训练的代码里设置了不同的batch_size，前者是8后者是4，有可能冻结训练的时候显存是够用的，结果解冻后显存不足了，这个时候需要重新把解冻训练阶段的batch_size调得更小一点。但是网络才训练了冻结阶段的50个epoch，backbone参数还是用的预训练权重呢，网络效果肯定不够好。难道要前功尽弃重新开始训练？这时候就要使用断点恢复技术了。其实断点恢复的思想很简单，就是把网络初始设置的model_path改为出错前保存好的权值文件，然后调整一下起始epoch和终止epoch即可，比如在前面提到的这种情况里，在第51个epoch报了错，那么可以把model_path修改为第50个epoch训练结束后保存的权值文件，然后把起始epoch调整成50就可以了。

断点恢复的应用范围非常非常广。最常见的情况就是代码跑到一半因为某些原因中断了（比如电脑突然死机重启这种不可抗力因素），又不想从头重新跑，那么就可以利用断点恢复训练的方法，这样可以节省不少时间。再比如，一个非常常见的情况，假如一开始设置了100个epoch，结果模型训练结束时，Loss还呈现下降的趋势，也就是模型还没有收敛，这种现象有可能就是epoch设置小了，所以可以把第100个epoch训练得到的权值文件当做初始权值文件再训练几个epoch看看，避免重新设置epoch从头训练。

当然，想要执行断点恢复首先需要把每个epoch得到的权值文件保存起来，这样才能修改model_path重新加载。断点恢复和常规的模型保存加载的区别其实就是epoch也要修改一下而已。保存权重可以用以下方法：

torch.save(model.state_dict(), "你要保存到的路径")

05

预训练和微调

最后再来总结一下预训练和微调，这是两个非常重要的概念，其实也很好理解。举个栗子是最能直观理解的。

假如我们现在要搭建一个网络模型来完成一个图像分类的任务，首先我们需要把网络的参数进行初始化，然后在训练网络的过程中不断对参数进行调整，直到网络的损失越来越小。在训练过程中，一开始初始化的参数会不断变化，如果结果已经满意了，那我们就可以把训练好的模型参数保存下来，以便训练好的模型可以在下次执行类似任务的时候获得比较好的效果。这个过程就是预训练（Pre-Training）。

假如在完成上面的模型训练后，我们又接到另一个类似的图像分类任务，这时我们就可以直接使用之前保存下来的模型参数作为这一次任务的初始化参数，然后在训练过程中依据结果不断进行修改，这个过程就是微调（Fine-Tuning）。

我们使用的神经网络越深，就需要越多的样本来进行训练，否则就很容易出现过拟合现象。比如我们想训练一个识别猫的模型，但是自己标注数据精力有限只标了100张，这时就可以考虑ImageNet数据集，可以在ImageNet上训练一个模型，然后使用该模型作为类似任务的初始化或特征提取器，这样既节省了时间和计算资源，又能很快地达到较好的效果。当然，采用预训练+微调也不是绝对有效的，上面识别猫的例子可以这样做是因为ImageNet里有猫的图像，所以可以认为是一个类似的数据集，如果是识别癌细胞的话，效果可能就不是那么好了。关于预训练和微调是有很多策略的，经验也很重要。

猜您喜欢：

超110篇！CVPR 2021最全GAN论文汇总梳理！

超100篇！CVPR 2020最全GAN论文梳理汇总！

拆解组新的GAN：解耦表征MixNMatch

StarGAN第2版：多域多样性图像生成

附下载 | 《可解释的机器学习》中文版

附下载 |《TensorFlow 2.0 深度学习算法实战》

附下载 |《计算机视觉中的数学方法》分享

《基于深度学习的表面缺陷检测方法综述》

《零样本图像分类综述: 十年进展》

《基于深度神经网络的少样本学习综述》