Pytorch：学习率设置笔记

torch.optim.lr_scheduler 包中封装了几种基于epoch的学习率调整方法。

一、StepLR

• torch.optim.lr_scheduler.StepLR(optimizer, step_size, gamma=0.1, last_epoch=-1)

Sets the learning rate of each parameter group to the initial lr decayed by gamma every step_size epochs. When last_epoch=-1, sets initial lr as lr.

# Assuming optimizer uses lr = 0.05 for all groups
# lr = 0.05     if epoch < 30
# lr = 0.005    if 30 <= epoch < 60
# lr = 0.0005   if 60 <= epoch < 90
# ...
scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
for epoch in range(100):
    scheduler.step()
    train(...)
    validate(...)

net = tv.models.resnet18()
optimizer = torch.optim.SGD(net.parameters(), lr=0.5, momentum=0.9)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
for epoch in range(10):
    scheduler.step()
    print(optimizer.state_dict()['param_groups'][0]['lr'])
    # train(...)
    # validate(...)

'''
0.5
0.5
0.5
0.05
0.05
0.05
0.005
0.005
0.005
0.0005
'''

 

二、MultiStepLR

• torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones, gamma=0.1, last_epoch=-1)

Set the learning rate of each parameter group to the initial lr decayed by gamma once the number of epoch reaches one of the milestones. When last_epoch=-1, sets initial lr as lr.

>>> # Assuming optimizer uses lr = 0.05 for all groups
>>> # lr = 0.05     if epoch < 30
>>> # lr = 0.005    if 30 <= epoch < 80
>>> # lr = 0.0005   if epoch >= 80
>>> scheduler = MultiStepLR(optimizer, milestones=[30,80], gamma=0.1)
>>> for epoch in range(100):
>>>     scheduler.step()
>>>     train(...)
>>>     validate(...)

三、其他学习率调整方法

此外Pytorch库中还有一些其他的学习率调整方法，如下：

• torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_epoch=-1)
• torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma, last_epoch=-1)
• torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max, eta_min=0, last_epoch=-1)
• torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=10, verbose=False, threshold=0.0001, threshold_mode='rel', cooldown=0, min_lr=0, eps=1e-08)

 

参考

• https://pytorch.org/docs/stable/optim.html