PyTorch中的动态学习率

PyTorch有两种动态调整学习率的方法，一种是使用PyTorch自带的 lr_scheduler，另一种是手动调整，可定制性更强。

一、使用 lr_scheduler

一个可能的例子：

import torch
import torch.nn as nn
from torch.optim import lr_scheduler

dataset = ...  # 数据集
model = ...  # 模型
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)  # 优化器
critertion = nn.CrossEntropyLoss()  # 损失函数
scheduler = lr_scheduler.ExponentialLR(optimizer, gamma=0.9)  # 学习率调度器

for epoch in range(20):
    for inputs, target in dataset:
		# 前向传播
        pred = model(inputs)
        loss = criterion(pred, target)
		
		# 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

	# 更新学习率
    scheduler.step()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

注意，scheduler.step() 需要在每个epoch结束后调用，且必须在 optimizer.step() 后调用。

更一般的模板：

for epoch in range(NUM_EPOCHS):
    train(...)  # 训练
    validate(...)  # 验证
    scheduler.step()  # 更新学习率
1
2
3
4

二、手动调整

有些时候我们需要手动调整学习率，例如Transformer的学习率定义如下

$$\text{lr}=d_{\text{model}}^{-0.5}\cdot \min (step\_nu{m}^{-0.5},step\_num\cdot warmup\_step{s}^{-1.5})$$

实现如下

class TransformerLR:
    def __init__(self, optimizer, d_model, warmup_steps=4000):
        self.optimizer = optimizer
        self.d_model = d_model
        self.warmup_steps = warmup_steps
        self.num_step = 0

        self.step()

    def step(self):
        new_lr = self.d_model**(-0.5) * min(self.num_step**(-0.5), self.num_step * self.warmup_steps**(-1.5))
        for group in self.optimizer.param_groups:
            group['lr'] = new_lr
        self.num_step += 1
1
2
3
4
5
6
7
8
9
10
11
12
13
14

在定义优化器时，我们可以任意设置初始学习率（一般设为0），这是因为 TransformerLR 在实例化时会自动调用 step 方法对优化器的学习率进行更新。

optimizer = torch.optim.Adam(model.parameters(), lr=0., betas=(0.9, 0.98), eps=1e-9)
scheduler = TransformerLR(optimizer, d_model=512)
1
2

---

References

[1] https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate
[2] https://www.ylkz.life/deeplearning/p10550146/
[3] http://www.4k8k.xyz/article/qq_36102055/119321243