Pytorch实现warm up和consine decay

news2024/11/24 4:05:59

在深度学习领域，模型训练过程中的不稳定性是一个常见的问题。为了解决这个问题，在Resnet这篇论文也提及了Warm Up的方法，通过逐渐增加学习率，引导模型在训练初期更稳定地收敛。同时在warm up之后结合consine decay的方法让训练变得更有效。

warm up和consine decay的意义

warm up来自于这篇文章：https://arxiv.org/pdf/1706.02677.pdf
consine decay来自于这篇文章：https://arxiv.org/pdf/1812.01187.pdf

在第一轮训练的时候，每个数据点对模型来说都是新的，模型会很快地进行数据分布修正，如果这时候学习率就很大，极有可能导致开始的时候就对该数据“过拟合”，后面要通过多轮训练才能拉回来，浪费时间。
当训练了一段时间（比如两轮、三轮）后，模型已经对每个数据点看过几遍了，或者说对当前的batch而言有了一些正确的了解，较大的学习率就不那么容易会使模型学偏，所以可以适当调大学习率。这个过程就可以看做是warmup。
那么为什么之后还要decay呢？当模型训到一定阶段后（比如10个epoch），模型的分布就已经比较固定了，或者说能学到的新东西就比较少了。如果还沿用较大的学习率，就会破坏这种稳定性，用我们通常的话说，就是已经接近loss的local optimal了，为了靠近这个最低点，我们就要慢慢来。

代码实现

1. 非常简单的前期准备工作

import torch,math
import matplotlib.pyplot as plt

#　定义一个简单的网络
class Model(torch.nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.linear = torch.nn.Linear(1, 1, bias=False)

    def forward(self, x):
        return self.linear(x)

# 声明一些超参数
epochs = 100000
warm_up_epochs = epochs*0.3
milestones = [epochs*0.3, epochs*0.7]

# 优化器
optimizer = torch.optim.SGD(model.parameters(), 0.1, momentum=0.9, weight_decay=5e-4)

2. 设置scheduler调度器


# MultiStepLR without warm up
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=milestones, gamma=0.1)

# warm_up_with_multistep_lr
warm_up_with_multistep_lr = lambda epoch: epoch / warm_up_epochs if epoch <= warm_up_epochs else 0.1**len([m for m in milestones if m <= epoch])
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=warm_up_with_multistep_lr)

# warm_up_with_cosine_lr
warm_up_with_cosine_lr = lambda epoch: epoch / warm_up_epochs if epoch <= warm_up_epochs else 0.5 * ( math.cos((epoch - warm_up_epochs) /(epochs - warm_up_epochs) * math.pi) + 1)
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=warm_up_with_cosine_lr)

上面的三段代码分别是：

不使用warm up + multistep learning rate 衰减
使用warm up + multistep learning rate 衰减
使用warm up + consine learning rate衰减

代码均使用pytorch中的torch.optim.lr_scheduler.LambdaLR自定义学习率衰减器

3. 模拟训练过程

lrs = []
for item in range(epochs):
    lr = optimizer.param_groups[0]["lr"]
    optimizer.zero_grad()
    optimizer.step()
    scheduler.step()
    lrs.append(scheduler.get_last_lr())
   
plt.plot(range(epochs),lrs)
plt.show()