一文搞懂pytorch中的学习率优化策略，torch.optim.lr_scheduler

一文搞懂pytorch中的学习率优化策略，torch.optim.lr_scheduler

torch.optim.lr_scheduler提供了三种根据epoch训练数来调整学习率（learning rate）的方法。具体方法如下：

一、有序调整：

1.等间隔调整(Step)：torch.optim.lr_scheduler.StepLR

$$ne{w}_{-}lr=initia{l}_{-}lr\ast {\gamma }^{\frac{epoch}{ste{p}_{-}size}}$$

更新策略：
每过step_size个epoch，学习率进行一次更新

$ne{w}_{-}lr$：新得到的学习率。
$initia{l}_{-}lr$：初始学习率。
$\gamma$：参数gamma，更新$lr$的乘法因子。
$ste{p}_{-}size$：每训练step_size个epoch，更新一次$lr$。
$las{t}_{-}epoch$：最后一个epoch的index，如果是训练了很多个epoch后中断了，继续训练，这个值就等于加载的模型的epoch。默认为-1表示从头开始训练，即从epoch=1开始。

代码如下：

import matplotlib.pyplot as plt
import torch
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR
from torchvision.models import AlexNet

num_epochs = 100

#定义2分类网络
model = AlexNet(num_classes=2)
optimizer = optim.SGD(params=model.parameters(), lr=0.05)

scheduler = StepLR(
    optimizer=optimizer,
    step_size=20,  # 设定调整的间隔数
    gamma=0.5,  # 系数
    last_epoch=-1
)
# 迭代训练
lrs, epochs = [], []
for epoch in range(num_epochs):
    lrs.append(scheduler.get_lr())  #.get_lr()获取当前学习率
    epochs.append(epoch)

    pass  # 在这里进行迭代训练
    #学习率更新
    scheduler.step()
    
# visualize
plt.figure()
plt.legend()
plt.plot(epochs, lrs, label='StepLR')
plt.show()

2.按需调整学习率(MultiStep)：

torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones, gamma=0.1, last_epoch=-1)
$$ne{w}_{-}lr=initia{l}_{-}lr\ast {\gamma }^{bisec{t}_{-}right(milestones,epoch)}$$
更新策略：
每次遇到milestones中的epoch，学习率做一次更新。

$ne{w}_{-}lr$：新得到的学习率。
$initia{l}_{-}lr$：初始学习率。
$\gamma$：参数gamma，更新$lr$的乘法因子。
$milestones$：递增的list，存放需要更新$lr$的epoch。
$las{t}_{-}epoch$：最后一个epoch的index，如果是训练了很多个epoch后中断了，继续训练，这个值就等于加载的模型的epoch。默认为-1表示从头开始训练，即从epoch=1开始。
$bisec{t}_{-}right(milestones,epoch)$：bisect模块中的bisect_right函数，返回值是把epoch插入排序好的列表milestones式的位置。

代码如下：

import matplotlib.pyplot as plt
import torch
import torch.optim as optim
from torch.optim.lr_scheduler import MultiStepLR
from torchvision.models import AlexNet

num_epochs = 100

#定义2分类网络
model = AlexNet(num_classes=2)
optimizer = optim.SGD(params=model.parameters(), lr=0.05)

scheduler = MultiStepLR(
    optimizer=optimizer,
    milestones=[10, 20, 40],  # 设定调整的间隔数
    gamma=0.5,  # 系数
    last_epoch=-1 
)

# train-like iteration
lrs, epochs = [], []
for epoch in range(num_epochs):
    lrs.append(scheduler.get_lr())  #.get_lr()获取当前学习率
    epochs.append(epoch)

    pass  # 在这里进行迭代训练
    #学习率更新
    scheduler.step()

# visualize
plt.figure()
plt.legend()
plt.plot(epochs, lrs, label='MultiStepLR')
plt.show()

其中最后一个epoch的index，如果是训练了很多个epoch后中断了，继续训练，这个值就等于加载的模型的epoch。默认为-1表示从头开始训练，即从epoch=1开始。

3.指数衰减调整(Exponential)：torch.optim.lr_scheduler.ExponentialLR
$$ne{w}_{-}lr=initia{l}_{-}lr\ast {\gamma }^{epoch}$$

import matplotlib.pyplot as plt
import torch
import torch.optim as optim
from torch.optim.lr_scheduler import ExponentialLR
from torchvision.models import AlexNet
import seaborn as sns

num_epochs = 100

#定义2分类网络
model = AlexNet(num_classes=2)
optimizer = optim.SGD(params=model.parameters(), lr=0.05)

scheduler = ExponentialLR(
    optimizer=optimizer,
    gamma=0.5,  # 系数
    last_epoch=-1
)

# train-like iteration
lrs, epochs = [], []
for epoch in range(num_epochs):
    lrs.append(scheduler.get_lr())  #.get_lr()获取当前学习率
    epochs.append(epoch)

    pass  # 在这里进行迭代训练
    #学习率更新
    scheduler.step()

# visualize
plt.figure()
plt.plot(epochs, lrs, label='ExponentialLR')
plt.legend()
plt.show()

4.余弦退火调整（CosineAnnealing）：torch.optim.lr_sheduler.CosineAnnealingLR
$$ne{w}_{-}lr=et{a}_{-}min+(initia{l}_{-}lr-et{a}_{-}min)\ast (1+cos(\frac{epoch}{T_{-}max}\ast \pi ))$$
更新策略：
每次遇到milestones中的epoch，学习率做一次更新。

$ne{w}_{-}lr$：新得到的学习率。
$initia{l}_{-}lr$：初始学习率。
$et{a}_{-}min$：表示最小学习率。
$T_{-}max$：代表1/2个cos周期所对应的epoch值，学习率下降到最小值时的epoch数，即当epoch=T_max时，学习率下降到余弦函数最小值，当epoch>T_max时，学习率将增大；
$las{t}_{-}epoch$：最后一个epoch的index，如果是训练了很多个epoch后中断了，继续训练，这个值就等于加载的模型的epoch。默认为-1表示从头开始训练，即从epoch=1开始。

import matplotlib.pyplot as plt
import torch
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR
from torchvision.models import AlexNet

num_epochs = 100

#定义2分类网络
model = AlexNet(num_classes=2)
optimizer = optim.SGD(params=model.parameters(), lr=0.05)

scheduler = CosineAnnealingLR(
    optimizer=optimizer,
    T_max=20,  # 设定调整的间隔数
    eta_min=0.005,  # 系数
    last_epoch=-1
)
# 迭代训练
lrs, epochs = [], []
for epoch in range(num_epochs):
    lrs.append(scheduler.get_lr())  #.get_lr()获取当前学习率
    epochs.append(epoch)

    pass  # 在这里进行迭代训练
    #学习率更新
    scheduler.step()
    
# visualize
plt.figure()
plt.legend()
plt.plot(epochs, lrs, label='CosineAnnealingLR')
plt.show()

二、自适应调整：
自适应调整学习率 ReduceLROnPlateau:torch.optim.lr_scheduler.ReduceLROnPlateau
$$ne{w}_{-}lr=\lambda \ast ol{d}_{-}lr$$

三、自定义调整：
自定义调整学习率 LambdaLR：torch.optim.lr_scheduler.LambdaLr(optimizer, lr_lambda, last_epoch=-1)

更新策略：
遇到模型不同的层参数时，根据该层参数的调整方式来调整学习率。

为不同参数组设定不同学习率调整策略，在fine-tune中特别有用，我们不仅可以为不同层设置不同的学习率，还可以为不同层设置不同的学习率调整策略。
$$ne{w}_{-}lr=\lambda \ast initia{l}_{-}lr$$
$ne{w}_{-}lr$：新得到的学习率。
$initia{l}_{-}lr$：初始学习率。
$\lambda$：通过参数lr_lambda和epoch得到的。
$l{r}_{-}lambda$：根据epoch计算λ的函数；或者是一个list的这样的function，分别计算各个parameter groups的学习率更新用到的λ；

import matplotlib.pyplot as plt
import torch
import torch.optim as optim
from torch.optim.lr_scheduler import LambdaLR
from torchvision.models import AlexNet
import seaborn as sns

num_epochs = 100

#定义2分类网络
model = AlexNet(num_classes=2)


# optimizer parameter groups 设置了个优化组：权重，偏置，其他参数
pg0, pg1, pg2 = [], [], [] 
for k, v in model.named_parameters():
    v.requires_grad = True
    if '.bias' in k:
        pg2.append(v)  # biases
    elif '.weight' in k and '.bn' not in k:
        pg1.append(v)  # apply weight decay
    else:
        pg0.append(v)  # all else

optimizer = optim.SGD(pg1, lr=0.05)

#给optimizer管理的参数组中增加新的组参数，
#可为该组参数定制lr,momentum,weight_decay 等在finetune 中常用。
optimizer.add_param_group({'params': pg2})  # add pg2 (biases)

# 对于卷积层权重的学习率，每十轮调整为原来的 0.1 倍
pg1 = lambda epoch: 0.1 ** (epoch // 10)
# 对于偏置权重的学习率，每轮调整为原来的 0.94 倍
pg2 = lambda epoch: 0.94 ** epoch


scheduler = LambdaLR(
    optimizer=optimizer,
    lr_lambda=[pg1, pg2],  #传入一个函数或一个以函数为元素列表，作为学习率调整的策略
    last_epoch=-1
)

weight_lrs, bias_lrs, epochs = [], [], []
for epoch in range(num_epochs):
    weight_lrs.append(scheduler.get_lr()[0])
    bias_lrs.append(scheduler.get_lr()[1])
    epochs.append(epoch)

    pass  # iter and train here

    scheduler.step(epoch)

# visualize
plt.figure()
plt.plot(epochs, weight_lrs, label='step')
plt.plot(epochs, bias_lrs, label='exponential')
plt.legend()
plt.show()