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 深度学习算法实战》

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

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

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

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

44d0ca7a798a1dc830d6310720c5f829.png

  • 19
    点赞
  • 113
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
### 回答1: PyTorch使用预训练权重的方法通常包括以下几个步骤: 1. 选择需要使用的预训练模型。PyTorch官方提供了许多流行的预训练模型,如ResNet、VGG、Inception等,可以在torchvision.models中找到。 2. 加载预训练权重。可以使用torchvision.models中的预训练模型的load_state_dict()方法来加载预训练权重,或者直接使用torch.load()方法加载已经保存好的预训练权重。 3. 将模型适配到自己的任务上。根据自己的任务需要,可以修改预训练模型的最后一层或几层,或者在预训练模型之上添加新的层。 4. 训练模型。使用加载了预训练权重的模型进行训练,通常需要在训练过程中冻结预训练模型的部分层参数,只对需要修改的层进行训练。 以上是使用预训练权重的一般流程,具体实现可以参考PyTorch官方文档中的示例代码。 ### 回答2: PyTorch是一种流行的深度学习框架,使我们能够方便地使用预训练权重进行模型初始化,进而提升模型的性能和加快训练速度。 预训练权重是在大规模数据集上训练好的模型参数。使用预训练权重的好处是,它们包含了在大量数据上学习到的有用特征,可以作为模型初始化的一种方式。 在PyTorch中,我们可以通过下载预训练权重来使用它们。例如,torchvision包提供了从ImageNet数据集预训练的许多常用模型,如ResNet、VGG、AlexNet等。通过使用这些预训练权重,我们可以获得在图像识别任务上具有很高准确性的模型。 使用预训练权重的步骤如下: 1. 首先,导入PyTorch和所需的预训练权重模型。 2. 创建模型实例,并加载预训练权重。 3. 将模型设置为评估模式,即不进行梯度计算。 4. 输入新的数据进行预测。 在加载预训练权重时,我们可以选择冻结一部分或全部权重冻结权重意味着它们在训练过程中将保持不变,只有其他参数会更新。这对于微调模型非常有用,即在新数据集上进行训练,以适应特定任务。 当我们使用预训练权重时,模型通常能够更快收敛,并且在训练集上获得更好的初始性能。然而,对于特定任务,预训练权重可能不一定是最佳选择。在某些情况下,我们可能需要进行微调或自定义的权重。 使用PyTorch预训练权重,我们能够方便地利用先前在大规模数据集上训练的模型参数,从而加速模型训练并提高模型性能。 ### 回答3: PyTorch是一个流行的深度学习框架,可以使用预训练模型权重来帮助我们快速建立和训练模型。预训练权重是在大规模数据集上预训练的模型参数,可以作为初始参数或微调参数用于特定任务。 使用预训练权重的第一步是选择适合任务的模型架构。PyTorch提供了一系列常用的预训练模型,如ResNet、VGG、AlexNet等。你可以根据任务的需求选择合适的模型。 下载预训练权重是接下来的一步。PyTorch通过torchvision提供了一种简便的方式来下载和加载预训练权重。你可以使用torchvision.models模块中的函数,如resnet50(pretrained=True)来下载ResNet-50的预训练权重。 加载预训练权重后,你可以将其应用于你的任务。如果你希望在预训练权重的基础上微调模型,则需要冻结一些层,只更新部分参数。通过设置requires_grad为False,可以冻结权重,使其不参与梯度更新。 一旦你完成了模型的设置和参数初始化,就可以开始训练了。你可以使用预训练权重作为模型的初始参数,让模型更快地收敛和获得更好的性能。 需要注意的是,预训练权重通常是在大规模数据集上训练得到的,因此可能适用于许多相关任务。然而,如果你的任务与预训练模型的训练数据具有很大的差异,可能需要进行微调或重新训练以适应你的任务。 总而言之,PyTorch提供了便捷的方式来使用预训练权重,通过应用预训练权重,我们可以更快地构建和训练模型,并获得更好的性能。

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值