一文看懂学习率warmup及各主流框架实现差异

最新推荐文章于 2024-07-29 09:33:03 发布

小白学视觉

最新推荐文章于 2024-07-29 09:33:03 发布

阅读量606

点赞数

文章标签：学习

原文链接：https://mp.weixin.qq.com/s?__biz=MzU0NjgzMDIxMQ==&mid=2247614590&idx=3&sn=412efd1e5b9c455cb66500a299429350&chksm=fabe10b2132efdff9f6cba5cf33025fc948fee89fb32a6080934e8c0b2f0d7f1b308753f55ae&scene=126&sessionid=0

版权

点击上方“小白学视觉”，选择加"星标"或“置顶”

重磅干货，第一时间送达

编者荐语

文章介绍了什么是warmup，几种常见的warmup方式，Timm、Detectron2和MMCV对于linear warmup的不同实现方式，在已有几个库实现方式下为啥还要自己写的原因，以及代码实现方式。

作者丨serendipity@知乎

链接丨https://zhuanlan.zhihu.com/p/508953700

什么是学习率warmup（热身）？

相关问题：神经网络中 warmup 策略为什么有效；有什么理论解释么？

目前还没有很严谨的理论证明，只有一些直觉上的解释，毕竟是炼丹。在网络训练前期，初始化的权重离最优权重距离较远，因此loss很大，计算出来的梯度也很大，此时如果使用较大的学习率，反而会让网络的优化过犹不及，距离最优点越来越远，通俗来说就是“跑飞了”。有些研究表明，训练初期跑飞了，后期不一定能拉的回来。

举个更加形象的例子，假设损失函数 l 是一个关于权重 w 的高阶函数（神经网络可以看成是一个高阶函数），，。显然， l(w) 的全局最优位置为 w=0 。假设初始位置为，我们知道一阶函数的梯度就是导数，因此根据SGD的更新原则可以得到 η，是更新之后的权重， η 是学习率。代入数值之后，得到 η ，可以看到如果学习率过大（ η=0.1 ），很容易跳过最优值，到达一个比当前更加差的位置（ w1=−17.2 ）。因此在训练初期，loss很大（或者说梯度很大）的时候，更适合使用较小的学习率。等到网络学会了一些知识，loss不那么大了，就可以换用更大的学习率，然后按照正常规则（MultiStep，CosineAnnealing）慢慢衰减。

几种常见的warmup方式

实验设置：使用的学习率衰减策略为MultiStep，milestones设置为[15, 20]，即学习率会在第16个epoch和第21个epoch衰减，乘以0.1的系数。整个训练过程持续了30个epoch，初始学习率 η0 为5，warmup_epochs设置为10，warmup_factor设置为0.001。

如果不使用warmup，学习率变化曲线如下，横坐标为epoch（从0开始索引），纵坐标为学习率。

multistep + no warmup

常见的warmup方式有三种：constant，linear和exponent。

constant：在warmup期间,学习率η=η0×warmup_factor=5×0.001=0.005。ResNet论文中就使用了这种方式，在cifar10上训练ResNet 101时，先用0.01的学习率训练直到训练误差低于80%（大概训练了400个steps），然后使用0.1的学习率进行训练。

multistep + constant warmup

linear：constant的不足之处在于从一个很小的学习率一下变为比较大的学习率可能会导致训练误差突然增大。linear方式可以避免这种问题，在warmup期间，学习率从η0×warmup_factor 线性增长到η0 。

multistep + linear warmup

exponent：在warmup期间，学习率从η0×warmup_factor指数增长到η0。

multistep + exponent warmup

Timm、Detectron2和MMCV对于linear warmup的不同实现方式

Timm（全称pytorch-image-models，by Ross Wightman）、Detectron2（by facebook）和MMCV（by SenseTime）是计算机视觉领域非常热门的开源代码库，然而这些库对于linear warmup的实现有着细微差异。核心代码片段放在这里以供验证，Timm、Detectron2、MMCV。

实验设置：使用的学习率衰减策略为CosineAnnealing。整个训练过程持续了30个epoch，初始学习率 η0 为5，warmup_epochs设置为10，warmup_factor设置为0.001（适用于Detectron2和MMCV），warmup_init_lr设置为0.005（适用于Timm）。

如果不使用warmup，学习率变化曲线LR(epoch) 如下。

cosine + no warmup

需要注意的是，在代码实现中warmup包含两个过程。

当0≤epoch<warmup_epochs 时，根据不同的策略设置相应的warmup学习率。
当epoch=warmup_epochs 时，将学习率设置为LR(warmup_epochs) ，这样可以保证在warmup结束之后学习率能够恢复到常规数值。

Timm：如果想让学习率在warmup_epochs个epoch内，线性地从warmup_init_lr增长到η0 ，每个epoch的增长量, 也就是斜率为Î· 。最终的学习率为warmup_init_lr (第0个epoch)、warmup_init_lr+slope (第1个epoch)、......、η0−slope (第warmup_epochs-1个epoch)、LR(warmup_epochs) (第warmup_epochs个epoch)。这里想要说明的是：η0 被用来计算斜率，并不意味着linear warmup的终点是η0 ，真正的终点是LR(warmup_epochs) 。一句话描述Timm的实现方式：不管最后的终点在哪，都先以η0 为目标线性地增长warmup_epochs个epoch。我将这种方式称呼为fix。fix容易导致第warmup_epochs个epoch的学习率出现跳变，如下图所示。

cosine + Timm linear warmup

Detectron2：和Timm不同，D2线性增长的目标并不是η0 ，而是LR(warmup_epochs) ，因此斜率的计算方式变为Î· 。这样可以避免学习率跳变的问题。在D2中初始学习率为η0×warmup_factor ，而不再是固定的warmup_init_lr 。这样的好处是可以为optimizer的不同param_groups自适应地设置初始学习率。我将这种方式称呼为auto。相较于fix，auto会自动地调整斜率，使得线性增长的终点刚好是LR(warmup_epochs) 。

cosine + Detectron2 linear warmup

MMCV：在warmup阶段 warmup_lr(epoch) = ratio(epoch) * LR(epoch) ，就是在常规学习率的基础之上乘以一个比例。这个比例是线性增长的，ratio(0) =warmup_factor , ratio(1) = warmup_factor + slope 、......、ratio(warmup_epochs-1)= 1- slope, ratio(warmup_epochs), 其中slope=(1−warmup_factor) /warmup_epochs 。虽然ratio是线性增长的，但是LR(epoch) 本身是非线性的，所以最终的乘积也是非线性的。我将这种方式称呼为factor。

cosine + MMCV linear warmup

既然已经有上述三个库，为什么还要造轮子

讲原因之前，需要先了解什么是by_epoch和warmup_by_epoch。by_epoch控制lr_scheduler应该在每个iteration结束后调用，还是在每个epoch结束后调用。PyTorch自带的lr_scheduler都是by_epoch的。warmup_by_epoch控制应该在每个iteration之后warmup，还是在每个epoch之后warmup。

1. by_epoch=True & warmup_by_epoch = True。Timm支持此模式。图中横坐标为iteration，纵坐标为学习率。注意之前的图中横坐标为epoch。

by_epoch=True & warmup_by_epoch=True

2. by_epoch=True & warmup_by_epoch = False，这种是笔者使用的最多的。MMCV支持此模式。

by_epoch=True & warmup_by_epoch=False

3. by_epoch=False & warmup_by_epoch = False，Timm、D2和MMCV均支持。

by_epoch=False & warmup_by_epoch=False

想造这个轮子主要有以下几个原因。

Timm、D2和MMCV都无法完全覆盖by_epoch和warmup_by_epoch的各种组合。
在Timm、D2和MMCV的代码设计中，需要自己再写一遍StepLR、MultiStepLR、CosineAnnealing。而这些都是PyTorch已经写好的，重复的代码完全可以避免。
MMCV的warmup相关代码被设计成了钩子，和MMCV.runner深度耦合，无法方便地拿出来作为一个单独的组件使用。
MMCV和D2均不支持ReduceLROnPlateau。
想写一个同时支持fix、auto、factor三种linear warmup方法的lr scheduler。

代码实现

本节将逐行讲解LRWarmupScheduler类的代码。这个类是我开源的Core-PyTorch-Utils(CPU) [https://github.com/serend1p1ty/core-pytorch-utils]库中的一个小组件，可以方便的放到其它项目中使用。

from torch.optim.lr_scheduler import ReduceLROnPlateau


class LRWarmupScheduler:
    """This class wraps the standard PyTorch LR scheduler to support warmup."""

    def __init__(self, torch_scheduler, by_epoch, epoch_len=None,
                 # 以下为warmup相关设置
                 warmup_t=0, warmup_by_epoch=False, warmup_mode="fix",
                 warmup_init_lr=None, warmup_factor=None):
        # PyTorch原生的lr scheduler
        self.torch_scheduler = torch_scheduler
        # 表示torch_scheduler是by epoch还是by iter
        self.by_epoch = by_epoch
        # 每个epoch的长度，只有by_epoch=True且warmup_by_epoch=False时才需要传入这个参数
        self.epoch_len = epoch_len
        # 后面有很多代码需要同时处理by epoch和by iter的情况，将变量命名为epoch或iter都不合适，
        # 所以选择用t来代表这层含义。当warmup_by_epoch=True时，warmup_t代表warmup_epochs，
        # 反之代表warmup_iters。
        self.warmup_t = warmup_t
        # 表示warmup是by epoch还是by iter
        self.warmup_by_epoch = warmup_by_epoch
        # 取值为fix、auto、factor
        self.warmup_mode = warmup_mode
        # fix模式的warmup初始学习率
        self.warmup_init_lr = warmup_init_lr
        # auto和factor模式下，warmup的初始学习率为base_lr * warmup_factor
        self.warmup_factor = warmup_factor

        self.param_groups = self.torch_scheduler.optimizer.param_groups
        self.base_lrs = [param_group["lr"] for param_group in self.param_groups]
        # 因为factor模式需要知道常规学习率才能推导出warmup学习率，所以需要预先计算出torch_scheduler在每个t的常规学习率。
        # 假设by_epoch=True & warmup_by_epoch=False & warmup_t=25 & epoch_len=10，
        # 说明warmup阶段跨越了3个epoch，我们需要预先计算出torch_scheduler在前三个epoch的
        # 常规学习率（保存在self.regular_lrs_per_t中）。
        # PS：虽然很多PyTorch原生的lr scheduler（StepLR、MultiStepLR、CosineAnnealingLR）
        # 提供了学习率的封闭形式，即_get_closed_form_lr()函数，可以通过传入的epoch参数直接计算出对应的学习率。
        # 但仍有些scheduler并未提供此功能，例如CosineAnnealingWarmRestarts。
        # 所以这里只能通过step()函数来一步步地计算出学习率。
        max_t = warmup_t // epoch_len if by_epoch and not warmup_by_epoch else warmup_t
        self.regular_lrs_per_t = self._pre_compute_regular_lrs_per_t(max_t)

        self.last_iter = self.last_epoch = 0
        self.in_iter_warmup = False

        if warmup_by_epoch:
            assert by_epoch
        if by_epoch and not warmup_by_epoch:
            assert epoch_len is not None
        if self._is_plateau:
            assert by_epoch
        if warmup_t > 0:
            if warmup_mode == "fix":
                assert isinstance(warmup_init_lr, float)
                # 为第0个t准备好学习率
                self._set_lrs(warmup_init_lr)
            elif warmup_mode == "factor":
                assert isinstance(warmup_factor, float)
                self._set_lrs([base_lr * warmup_factor for base_lr in self.base_lrs])
            elif warmup_mode == "auto":
                assert isinstance(warmup_factor, float)
                self.warmup_end_lrs = self.regular_lrs_per_t[-1]
                self._set_lrs([base_lr * warmup_factor for base_lr in self.base_lrs])
            else:
                raise ValueError(f"Invalid warmup mode: {warmup_mode}")

    @property
    def _is_plateau(self):
        return isinstance(self.torch_scheduler, ReduceLROnPlateau)

    def _pre_compute_regular_lrs_per_t(self, max_t):
        regular_lrs_per_t = [self.base_lrs]
        if self._is_plateau:
            return regular_lrs_per_t * (max_t + 1)
        for _ in range(max_t):
            self.torch_scheduler.step()
            regular_lrs_per_t.append([param_group["lr"] for param_group in self.param_groups])
        return regular_lrs_per_t

    def _get_warmup_lrs(self, t, regular_lrs):
        # 为了简单，不再计算斜率，而是通过线性插值的方式获得warmup lr
        alpha = t / self.warmup_t
        if self.warmup_mode == "fix":
            return [self.warmup_init_lr * (1 - alpha) + base_lr * alpha for base_lr in self.base_lrs]
        elif self.warmup_mode == "factor":
            factor = self.warmup_factor * (1 - alpha) + alpha
            return [lr * factor for lr in regular_lrs]
        else:
            return [
                base_lr * self.warmup_factor * (1 - alpha) + end_lr * alpha
                for base_lr, end_lr in zip(self.base_lrs, self.warmup_end_lrs)
            ]

    def _set_lrs(self, lrs):
        if not isinstance(lrs, (list, tuple)):
            lrs = [lrs] * len(self.param_groups)
        for param_group, lr in zip(self.param_groups, lrs):
            param_group['lr'] = lr

    def epoch_update(self, metric=None):
        if not self.by_epoch:
            return

        self.last_epoch += 1
        if self.warmup_by_epoch and self.last_epoch < self.warmup_t:
            # 0 <= t < warmup_t时，根据不同的策略设置相应的warmup学习率
            self._set_lrs(self._get_warmup_lrs(self.last_epoch, self.regular_lrs_per_t[self.last_epoch]))
        elif self.warmup_by_epoch and self.last_epoch == self.warmup_t:
            # t == warmup_t时，将学习率恢复为常规学习率
            self._set_lrs(self.regular_lrs_per_t[-1])
        # in_iter_warmup=True时代表正在进行by iter的warmup，
        # lr已经被设置好了，此时torch_scheduler不能再执行step()函数修改lr
        elif not self.in_iter_warmup:
            if self._is_plateau:
                self.torch_scheduler.step(metric)
            else:
                self.torch_scheduler.step()

    def iter_update(self):
        if self.warmup_by_epoch:
            return

        self.last_iter += 1
        if self.last_iter < self.warmup_t:
            self.in_iter_warmup = True
            t = self.last_iter // self.epoch_len if self.by_epoch else self.last_iter
            self._set_lrs(self._get_warmup_lrs(self.last_iter, self.regular_lrs_per_t[t]))
        elif self.last_iter == self.warmup_t:
            self._set_lrs(self.regular_lrs_per_t[-1])
        else:
            # warmup结束之后，将in_iter_warmup变量置为False，此时torch_scheduler才可以进行正常的step()
            self.in_iter_warmup = False
            if not self.by_epoch:
                self.torch_scheduler.step()

    def state_dict(self):
        state = {key: value for key, value in self.__dict__.items() if key != "torch_scheduler"}
        state["torch_scheduler"] = self.torch_scheduler.state_dict()
        return state

    def load_state_dict(self, state_dict):
        self.torch_scheduler.load_state_dict(state_dict.pop("torch_scheduler"))
        self.__dict__.update(state_dict)

下载1：OpenCV-Contrib扩展模块中文版教程

在「小白学视觉」公众号后台回复：扩展模块中文教程，即可下载全网第一份OpenCV扩展模块教程中文版，涵盖扩展模块安装、SFM算法、立体视觉、目标跟踪、生物视觉、超分辨率处理等二十多章内容。


下载2：Python视觉实战项目52讲
在「小白学视觉」公众号后台回复：Python视觉实战项目，即可下载包括图像分割、口罩检测、车道线检测、车辆计数、添加眼线、车牌识别、字符识别、情绪检测、文本内容提取、面部识别等31个视觉实战项目，助力快速学校计算机视觉。


下载3：OpenCV实战项目20讲
在「小白学视觉」公众号后台回复：OpenCV实战项目20讲，即可下载含有20个基于OpenCV实现20个实战项目，实现OpenCV学习进阶。


交流群

欢迎加入公众号读者群一起和同行交流，目前有SLAM、三维视觉、传感器、自动驾驶、计算摄影、检测、分割、识别、医学影像、GAN、算法竞赛等微信群（以后会逐渐细分），请扫描下面微信号加群，备注：”昵称+学校/公司+研究方向“，例如：”张三 + 上海交大 + 视觉SLAM“。请按照格式备注，否则不予通过。添加成功后会根据研究方向邀请进入相关微信群。请勿在群内发送广告，否则会请出群，谢谢理解~

小白学视觉

关注

0
点赞
踩
1

收藏

觉得还不错? 一键收藏
0
评论
一文看懂学习率warmup及各主流框架实现差异

点击上方“小白学视觉”，选择加"星标"或“置顶”重磅干货，第一时间送达编者荐语文章介绍了什么是warmup，几种常见的warmup方式，Timm、Detectron2和MMCV对于linear warmup的不同实现方式，在已有几个库实现方式下为啥还要自己写的原因，以及代码实现方式。作者丨serendipity@知乎链接丨https://zhuanlan.zhihu.com/p/50895370...
复制链接

扫一扫