模型训练的学习率衰减策略之Pytorch实现余弦退火策略

在训练深度学习模型时，学习率是一个非常重要的超参数，它控制着参数更新的步长。过高或过低的学习率都可能导致训练过程不稳定或收敛速度过慢。

预热（Warmup）阶段

在训练初期，模型参数通常是随机初始化的，这意味着它们可能离最优解非常远。如果此时使用较大的学习率，可能会导致参数在更新过程中产生剧烈的振荡，从而无法稳定地接近最优解。预热阶段的目的是在训练的最初几个周期（epochs）内逐渐增加学习率。这样做的好处是：

• 减少梯度爆炸和消失的风险：较小的初始学习率有助于模型在训练开始时稳定地更新参数。
• 加速收敛：随着学习率的逐渐增加，模型能够以更快的速度接近最优解。

 余弦退火（Cosine Annealing）阶段

在预热阶段之后，模型逐渐进入稳定的学习阶段。此时，如果继续使用固定的学习率，可能会因为学习率过高而导致模型在最优解附近振荡，或者因为学习率过低而导致收敛速度过慢。余弦退火策略通过周期性地调整学习率，使其按照余弦函数的形式先下降后上升，再下降，从而帮助模型更精细地调整参数。这种策略的好处包括：

• 避免陷入局部最优解：周期性的学习率变化有助于模型跳出局部最优解，探索更大的解空间。
• 提高收敛速度：在下降阶段，学习率的逐渐降低有助于模型更精细地调整参数，从而更快地收敛到最优解。

为什么这么训练有助于模型更快收敛

结合预热和余弦退火策略的学习率调度器，能够在训练的不同阶段提供适当的学习率，从而加速模型的收敛过程。具体来说：

• 预热阶段通过逐渐增加学习率，帮助模型在训练初期稳定地更新参数，避免梯度爆炸和消失的风险，为后续的快速收敛打下基础。
• 余弦退火阶段通过周期性地调整学习率，使模型在收敛过程中既能保持一定的更新速度，又能避免在最优解附近振荡，从而更快地达到较高的训练精度。

import matplotlib.pyplot as plt
import torch.optim as optim
from torchvision.models import resnet18
from math import pi, cos
from torch.optim.optimizer import Optimizer


class CosineWarmupLr(object):
    """Cosine lr decay function with warmup.
    Lr warmup is proposed by `
        Accurate, Large Minibatch SGD:Training ImageNet in 1 Hour`
        `https://arxiv.org/pdf/1706.02677.pdf`
    Cosine decay is proposed by `
        Stochastic Gradient Descent with Warm Restarts`
        `https://arxiv.org/abs/1608.03983`
    Args:
        optimizer (Optimizer): optimizer of a model.
        batches (int): batches of one epoch.
        max_epochs (int): max_epochs to train.
        base_lr (float): init lr.
        final_lr (float): minimum(final) lr.
        warmup_epochs (int): warmup max_epochs before cosine decay.
        warmup_init_lr (float): warmup starting lr.
        last_iter (int): init iteration.
    Attributes:
        niters (int): number of iterations of all max_epochs.
        warmup_iters (int): number of iterations of all warmup max_epochs.
    """
    def __init__(self, optimizer, batches, max_epochs, base_lr, final_lr=0,
                 warmup_epochs=0, warmup_init_lr=0, last_iter=-1):
        if not isinstance(optimizer, Optimizer):
            raise TypeError('{} is not an Optimizer'.format(type(optimizer).__name__))
        self.optimizer = optimizer
        if last_iter == -1:
            for group in optimizer.param_groups:
                group.setdefault('initial_lr', group['lr'])
            last_iter = 0
        else:
            for i, group in enumerate(optimizer.param_groups):
                if 'initial_lr' not in group:
                    raise KeyError("param 'initial_lr' is not specified "
                                   "in param_groups[{}] when resuming an optimizer".format(i))

        self.baselr = base_lr
        self.learning_rate = base_lr
        self.niters = max_epochs * batches
        self.targetlr = final_lr
        self.warmup_iters = batches * warmup_epochs
        self.warmup_init_lr = warmup_init_lr
        self.last_iter = last_iter
        self.step()

    def get_lr(self):
        if self.last_iter < self.warmup_iters:
            self.learning_rate = self.warmup_init_lr + \
                (self.baselr - self.warmup_init_lr) * self.last_iter / self.warmup_iters   #初始预热学习率到基础学习率之间的差值，然后按照当前迭代次数占预热迭代次数的比例进行缩放，最后将这个缩放后的差值加到初始预热学习率上，得到当前的学习率。
        else:
            self.learning_rate = self.targetlr + (self.baselr - self.targetlr) * \
                (1 + cos(pi * (self.last_iter - self.warmup_iters) /
                         (self.niters - self.warmup_iters))) / 2

    def step(self, iteration=None):
        """Update status of lr.
        Args:
            iteration(int, optional): now training iteration of all max_epochs.
                Normally need not to set it manually.
        """
        if iteration is None:
            iteration = self.last_iter + 1
        self.last_iter = iteration
        self.get_lr()
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = self.learning_rate


if __name__ == '__main__':
    
    warmup_epochs = 5  # 
    max_epoch = 200  # 共有120个epoch，则用于cosine rate的一共有100个epoch
    lr_init = 0.1
    lr_final = 1e-5
    lr_warmup_init = 0.

    iter_per_epoch = 100

    model = resnet18(num_classes=10)
    optimizer = optim.SGD(model.parameters(), lr=lr_init, momentum=0.9)
    scheduler = CosineWarmupLr(optimizer, batches=iter_per_epoch, max_epochs=max_epoch, base_lr=0.1,
                               final_lr=lr_final, warmup_epochs=warmup_epochs, warmup_init_lr=lr_warmup_init)
    index = 0
    x = []
    y = []
    for epoch in range(max_epoch):
        for iter in range(iter_per_epoch):
            lr_c = optimizer.param_groups[0]['lr']
            y.append(lr_c)
            print(lr_c)
            scheduler.step()
            x.append(index)
            index += 1

    plt.figure(figsize=(10, 8))
    plt.xlabel('iteration')
    plt.ylabel('cosine rate')
    plt.plot(x, y, color='r', linewidth=2.0, label='cosine rate')
    plt.legend(loc='best')
    plt.show()

 目前主流的学习率动态调整策略有4种：1、指数衰减  2、固定步长衰减  3、多步长衰减  4、余弦退火衰减
具体参考：pytorch必须掌握的的4种学习率衰减策略 - 知乎 (zhihu.com)
https://zhuanlan.zhihu.com/p/93624972