Transformers中的动态学习率

1. 基线模型

我这里用的是当时学 Pytorch 时搭的简单模型¹，网络结构的示意图大致如下所示，是 FashionMNIST 的多分类任务

贴一下主要代码：

import os
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision import datasets
import time
from transformers import optimization
from torch.utils.tensorboard import SummaryWriter
import datetime

all_start = time.time()

os.environ['CUDA_VISIBLE_DEVICES'] = '0'
# 配置GPU

# 配置其他超参数
batch_size = 256
# 批次大小
# num_workers = 4  
num_workers = 0
# 对于Windows用户，这里应设置为0，否则会出现多线程错误
lr = 1e-4
# 学习率
epochs = 1
# 轮数

print("hyperParameter init")

image_size = 28
# 图片大小
data_transform = transforms.Compose([
    # transforms.ToPILImage(),   
    # 这一步取决于后续的数据读取方式，如果使用内置数据集则不需要
    transforms.Resize(image_size),
    transforms.ToTensor()
])
# 数据转换方式

train_data = datasets.FashionMNIST(
    root='./',
    train=True,
    download=True,
    transform=data_transform
)
# 使用torchvision自带数据集

print("dataset init")

train_loader = DataLoader(
    train_data,
    batch_size=batch_size,
    shuffle=True,
    num_workers=num_workers,
    drop_last=True
)
# 配置数据加载方式

print("dataloader init")

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 32, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, stride=2),
            nn.Dropout(0.3),
            nn.Conv2d(32, 64, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, stride=2),
            nn.Dropout(0.3)
        )
        self.fc = nn.Sequential(
            nn.Linear(64*4*4, 512),
            nn.ReLU(),
            nn.Linear(512, 10)
        )
        
    def forward(self, x):
        x = self.conv(x)
        x = x.view(-1, 64*4*4)
        x = self.fc(x)
        # x = nn.functional.normalize(x)
        return x

model = Net()
# 实例化
model = model.cuda()
# 加载到显卡

print("model init")

criterion = nn.CrossEntropyLoss()
# 实例化损失函数

steps = len(train_loader.dataset)/batch_size

optimizer = optim.Adam(model.parameters(), lr=1.0)
# 配置优化器

# scheduler = optimization.get_constant_schedule(optimizer, last_epoch=-1)

# scheduler = optimization.get_constant_schedule_with_warmup(optimizer, num_warmup_steps=100)

# scheduler = optimization.get_linear_schedule_with_warmup(
#     optimizer,
#     num_warmup_steps=100,
#     num_training_steps=steps
# )

# scheduler = optimization.get_polynomial_decay_schedule_with_warmup(
#     optimizer,
#     num_warmup_steps=100,
#     num_training_steps=steps,
#     lr_end = 1e-7,
#     power=3
# )

# scheduler = optimization.get_cosine_schedule_with_warmup(
#     optimizer,
#     num_warmup_steps=100,
#     num_training_steps=steps,
#     num_cycles=2
# )

scheduler = optimization.get_cosine_with_hard_restarts_schedule_with_warmup(
    optimizer,
    num_warmup_steps=100,
    num_training_steps=steps,
    num_cycles=2
)

# 该数据集一共234个step,每个batch有256个样本
# 配置学习率的控制器

def train(epoch):
    model.train()
    # 标明训练模式
    train_loss = 0
    log_dir=os.path.join('logs',datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))
    writer = SummaryWriter(log_dir)
    st_time = time.time()
    for batch, (data, label) in enumerate(train_loader):
        data, label = data.cuda(), label.cuda()
        # 加载训练数据
        optimizer.zero_grad()
        # 梯度清零
        output = model(data)
        # 前向传播
        loss = criterion(output, label)
        # 计算损失
        loss.backward()
        # 反向传播
        optimizer.step()
        # 更新参数
        scheduler.step()
        # 更新学习率
        train_loss += loss.item()*data.size(0)
        # 叠加loss
        ls_lr = scheduler.get_last_lr()
        print(ls_lr)
        writer.add_scalar("lr", ls_lr[0], batch)
        # 记录该step的学习率,返回的是列表格式,只有一个元素
        writer.add_scalar("loss", train_loss/((batch+1)*batch_size), batch)
        # 记录样本的平均loss
        print("finish batch {}".format(batch+1))
        # 打印进度信息
    writer.close()
    # 关闭书写器
    train_loss = train_loss/len(train_loader.dataset)
    # 该epoch平均loss
    print('Epoch:{} Training Loss:{:.6f} cost:{:.2f}s'.format(
        epoch, train_loss, time.time()-st_time))
    # 打印信息

print("train start")

for epoch in range(1, epochs+1):
    train(epoch)
    
print("train end")

print("total cost {:.2f}s".format(time.time()-all_start))

噢，对了，我用的 GPU，设备信息如下：

(torch) PS D:\Project\AmarToy> nvidia-smi
Sun Jun 19 20:23:56 2022       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 419.72       Driver Version: 419.72       CUDA Version: 10.1     |
|-------------------------------+----------------------+----------------------+
| GPU  Name            TCC/WDDM | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  GeForce MX150      WDDM  | 00000000:01:00.0 Off |                  N/A |
| N/A   44C    P8    N/A /  N/A |     64MiB /  2048MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

嗯，对，我还用了 TensorBoard，当时想看看 Loss 和学习率。

因为感觉不太对，这个东西用在 CV 上就很奇怪

刚开始的时候忘记分文件夹了，每次训练的展示数据都混在一起，画出来的图群魔乱舞

后来按时间建文件夹，每个 epoch 一个，然后把路径传入 SummaryWriter

我记得在 TF2 里面是传入 callbacks，那个好像是自动记录所要展示的变量

2. 动态学习率

这个就主要是参考博客²，也可以去看知乎的版本³，这个带目录

2.1 constant

这个没啥说的，恒定值，不清楚有啥用，我感觉直接在优化器里面指定就好了，在这里好像多此一举

具体的调用方式为：

optimizer = optim.Adam(model.parameters(), lr=1.0)
scheduler = optimization.get_constant_schedule(optimizer, last_epoch=-1)

这里的 last_epoch 参数指的是学习率从哪一个 step 开始更新，-1 的话是从头开始

因为某些学习率是与 step 相关的，所以如果我们想继续这些模型的训练，就要记住上次训练到哪个 step，从而获得相应的学习率

然后在每一个批次训练的结尾，在反向传播之后，更新一下：

optimizer.step()
# 更新参数
scheduler.step()
# 更新学习率

可视化的结果如下

这个 Loss 的结果还行，还算正常，其他的没啥

实现方式上，恒定返回为 1，或者其他定值

2.2 constant_with_warmup

其调用方式为：

scheduler = optimization.get_constant_schedule_with_warmup(optimizer, num_warmup_steps=100)

其中热身步数为 100，我们在这个 256 的 batch_size 设定之下，一共有 234 个 batch

来看看可视化的结果，这个 Loss 跟鬼一样，后面还有更花哨的

我们可以看到，前 100 步的学习率是稳定线性上升的，起点为 0，终点为我们此前设定好的 1.0

也就是说，在指定的热身步数以内，学习率缓缓上升到指定数值，之后保持不变

其实现方式也很简单，就是根据当前的 step 来进行判断

如果还在热身步数以内，就返回当前步数与热身步数的比值与设定学习率的乘积

如果已经过了热身步数，那就要返回我们此前指定的学习率

这里还有一个细节，要防止热身步数为 0，那可能会使得分母为 0

这个学习率的变化是与步数相关的，但是我们并没有在参数当中看到有关当前 step 的东西

这个好像是因为其内部有一个独立的 step 计数部分，每更新一次就记录一步

当然好像也可以选择传入 step 来获得该 step 对应的学习率

2.3 linear_with_warmup

其调用如下

scheduler = optimization.get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=100,
    num_training_steps=steps
)

热身步数同样设定为 100，还需要传入训练步数，这个我们在前面的代码计算出来了

来看看可视化结果，这 Loss 神头鬼脸的：

我们可以看到，前 100 步的学习率稳步上升，后面的学习率稳步下降

那么从实现方式上来说，热身期的学习率就是，当前步数与热身步数的比值与设定学习率的乘积

后面的学习率，应该是，总步数减去当前步数与总步数减去热身步数的比值，再乘上设定的学习率

我们对后半部分的学习率展开讲一下，也就是，其需要在训练步数减去热身步数这个长度之内归零

每次衰减的部分，应该是，当前步数减去热身步数，那么其留存部分就是训练步数减去当前步数

我们用留存部分除以总的长度，就是当前学习率的系数，可以见示意图：

当然需要注意的问题还是，比值的分母不能为 0，以及这个学习率需要保证为正值

2.4 polynomial_with_warmup

其调用方式如下：

scheduler = optimization.get_polynomial_decay_schedule_with_warmup(
    optimizer,
    num_warmup_steps=100,
    num_training_steps=steps,
    lr_end = 1e-7,
    power=3
)

同样设定热身步数为 100，还需要传入训练步数，这个我们在前面的代码计算出来了

其他参数是，学习率的最小值，以及多项式的次数

可视化结果如下，Loss 的趋势和上面差不多，先降后升再平稳：

这个大概算三个部分，热身、衰减，以及稳定

热身跟前面一样，线性递增，把比值作为系数就行

衰减部分就复杂一点了，涉及多项式，主要思想跟上面其实也差不多

我们在上一部分可以看到，学习率需要在热身之后、训练步长之前，归零

那么对于这一部分，就是在热身之后、训练步长之前，降到指定值

对于线性衰减，我们可以试着写个公式：

衰减部分是

$decay=current\_step-warmup\_step$

留存部分是

$remain=training\_step-current\_step$

总的长度是

$total=training\_step-warmup\_step$

所以比值是

$ratio=\frac{remain}{total}$

也就是线性衰减的系数

但是我们这里不一样，我们这里是降到指定值，所以要变化一下：

我们设的起点值是

$lr\_begin$

我们设的终点值是

$lr\_end$

所以说衰减的范围是

$lr\_range=lr\_begin-lr\_end$

对于上面的线性衰减，改写成这种格式就是

$current\_lr=lr\_range\ast ratio+lr\_end$

那么将其变换为系数的话，还需要除以起始值，也就是

$lr\_lambda=\frac{current\_lr}{lr\_init}$

由于$lr\_range=lr\_begin$，且$lr\_end=0$

所以说$lr\_lambda=ratio$

对于多项式衰减，其实也就是给系数加上相应的次方，即

$current\_lr=lr\_range\ast coe{f}^{power}+lr\_end$

然后再变换为学习率的系数

$lr\_lambda=\frac{current\_lr}{lr\_init}$

如果多项式的次数为 1，那就是线性衰减

以上说了热身和衰减两个部分，最后其实还有一个稳定

这个大概是因为，学习率没有彻底衰减到 0，所以可以继续

也就是说可以超过设定的训练步长，继续训练

2.5 cosine_with_warmup

其调用方式为：

scheduler = optimization.get_cosine_schedule_with_warmup(
    optimizer,
    num_warmup_steps=100,
    num_training_steps=steps,
    num_cycles=2
)

热身还是 100 步，然后 cosine 的循环次数为 2

可视化结果如下，Loss 先升后降：

奇怪啊，这是抽什么风，Loss 咋先升高呢

学习率的热身部分跟上面一致啊，都是线性递增的

我想到了第二个，带热身的常数，那个在热身期间也有升高 Loss

可能是样本的问题吧，感觉应该是样本顺序的问题

对于衰减部分，公式如下：

$lr\_lambda=\frac{1}{2}\ast (cos(2\pi \ast num\_cycles\ast ratio)+1)$

我们先简单写一下原始的公式

$y=\frac{1}{2}\ast (cos(2\pi x)+1)$

对于余弦函数而言，我们先用 +1 将 y 值的范围从 [-1,1] 变换到 [0,2]

再用系数 0.5 将 y 值压缩到 [0,1]，使其变为比值

那么对于内部，我们加了系数 2*pi

原余弦的几个关键点有 0，0.5pi，pi，1.5pi，2*pi

现在其变为 0，0.25，0.5，0.75，1.0，也就是留存部分占比的主要部分

由红变蓝，如下图所示：

但是我们还需要注意的一个点是

在$lr\_lambda$当中，$x$部分不止有$ratio$，还有$num\_cycles$

这也是一个放缩，使得函数在总的长度$total$之内经历$num\_cycles$个周期

由蓝变红，如下图所示

2.6 cosine_with_restarts_and_warmup

其调用方式如下：

scheduler = optimization.get_cosine_with_hard_restarts_schedule_with_warmup(
    optimizer,
    num_warmup_steps=100,
    num_training_steps=steps,
    num_cycles=2
)

100 步热身，2 次循环

可视化的结果如下，Loss 先升再降：

热身部分没啥好说的，线性递增，衰减部分与上面相似，公式如下：

$lr\_lambda=\frac{1}{2}\ast (cos(\pi \ast num\_cycles\ast ratio\%1.0)+1)$

同样，先大概写一个原始公式

$y=\frac{1}{2}\ast (cos(\pi x\%1.0)+1)$

相对于余弦函数，我们同样地将 y 变换到 [0,1]

对于内部，我们加了系数 pi

将其关键点变为0，0.5，1.0，1.5，2.0，也就是留存部分占比的主要部分

由红变蓝，如下图所示

需要注意的是，我们在这里还 mod 了 1

这也就意味着我们将 x 的变化范围圈定到 [0,1]，只要衰减的这部分

这里的$lr\_lambda$当中，$x$部分有$ratio$和$num\_cycles$，效果同上

所谓 restarts 指的就是我们只留下了余弦函数当中衰减的半边

所以当完成一个循环，切换到另一个循环的时候，学习率会突然僵硬地跳变

我们在这这里的最小值是 0 ，所以说当前步数超出训练步数的时候，学习率为 0

3. 自定义学习率

3.1 LambdaLR

如果我们要手动实现学习率的话，可以导入相关的包来构造函数

例如上文的 constant_with_warmup 可以是这样：

from torch.optim.lr_scheduler import LambdaLR
def get_constant_schedule_with_warmup(Optimizer, num_warmup_steps,last_epoch = -1):
    def lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1.0, num_warmup_steps))
        return 1.0
    return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch)

我们可以看到，当定义好学习率随步数的变化函数之后，我们使用 LambdaLR 来完成其他工作

LambdaLR 主要有 get_lr() 和 step() 这两个方法，我们先来说说 step()，这个用于更新学习率

step() 的主要思想是，每当训练前进一步，就更新一步学习率

这里面也要分情况，我们可以选择是否指定步数

如果不指定的话，就是从第 0 步开始更新，随后函数内部自动记录步数，用于返回后面的学习率

如果我们指定步数的话，相当于说在此时从这一步的下一步开始更新学习率，只是当前的训练可能并不是这一步的下一步

那么对于 get_lr() 而言，就是拿到基础的学习率，及其变化的系数，返回当前的学习率变化结果

具体的情况可以参照原博客³

上文提到，我们可以指定学习率从某一步开始变化，这个可以用来恢复模型的训练

也就是说，当我们中断训练后想要继续恢复，可以指定学习率更新的步数为中断时的步数

3.2 transformer

实现方式如下：

from torch.optim.lr_scheduler import LambdaLR
def get_customized_schedule_with_warmup(optimizer, num_warmup_steps, d_model=1.0, last_epoch=-1):
    def lr_lambda(current_step):
        current_step += 1
        arg1 = current_step ** -0.5
        arg2 = current_step * (num_warmup_steps ** -1.5)
        return (d_model ** -0.5) * min(arg1, arg2)
    return LambdaLR(optimizer, lr_lambda, last_epoch)

scheduler = get_customized_schedule_with_warmup(
    optimizer,
    num_warmup_steps=100,
    d_model=728
)

需要注意的是 get_customized_schedule_with_warmup() 函数的返回值

我们在这里设置的学习率都是 1.0，所以直接返回计算结果

或者也可以除以初始学习率，使其更为通用一些

可视化结果如下，这个 Loss 还算正常：

这个学习率的变化挺神奇的

首先是一个基础部分：

$base$

$d\_model\ast \ast -0.5$

感觉是先开平方，再做倒数

$\frac{1}{\sqrt{d\_model}}$

然后是另一个部分，需要取两者之间的最小：

$arg1$

$current\_step\ast \ast -0.5$

$\frac{1}{\sqrt{current\_step}}$

以及：

$arg2$

$current\_step\ast (num\_warm\_steps\ast \ast -1.5)$

$\frac{current\_step}{num\_warm\_step{s}^{1.5}}$

$\frac{current\_step}{{}^2\sqrt{num\_warm\_step{s}^3}}$

$\frac{current\_step}{num\_warm\_steps\sqrt{num\_warm\_steps}}$

$\frac{current\_step\sqrt{current\_step}}{num\_warm\_steps\sqrt{num\_warm\_steps}}\ast \frac{1}{\sqrt{current\_step}}$

$\frac{\sqrt{current\_ste{p}^3}}{\sqrt{num\_warm\_step{s}^3}}\ast \frac{1}{\sqrt{current\_step}}$

$\sqrt{\frac{current\_ste{p}^3}{num\_warm\_step{s}^3}}\ast \frac{1}{\sqrt{current\_step}}$

$\sqrt{(\frac{current\_step}{num\_warm\_steps}{)}^3}\ast \frac{1}{\sqrt{current\_step}}$

所以我们可以看到

$arg2=\sqrt{(\frac{current\_step}{num\_warm\_steps}{)}^3}\ast arg1$

也就是看当前步数是否到达了热身步数

如未到达，取较小的$arg2$；已到达则取$arg1$

然后再将这个部分，乘上我们之前的基础部分

$lr\_lambda=base\ast arg2whencurrent\_step<num\_warm\_steps$

$lr\_lambda=base\ast arg1whencurrent\_step>num\_warm\_steps$

---

1. 深入浅出PyTorch ↩︎

2. Transformers自定义学习了动态调整 ↩︎

3. Transformers之自定义学习率动态调整 ↩︎ ↩︎