官方文档:https://pytorch.org/docs/1.10.1/generated/torch.optim.lr_scheduler.LambdaLR.html
在python中,有个东西叫做匿名函数(lambda表达式),能够用于很方便的定义各种规则,这个LambdaLR也就可以理解成自定义规则去调整网络的学习率。从另一个角度理解,数学中的 λ \lambda λ一般是作为系数使用,因此这个学习率调度器的作用就是将初始学习率乘以人工规则所生成的系数 λ \lambda λ。
函数结构如下:
torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_epoch=-1, verbose=False)
参数:
一个例子如下。考虑epoch从0算起,比如我们想每3个epoch(即在第2,5,8个epoch结束后)将学习率减半,代码如下:
import torch
from torch import nn
import math
class Net(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(in_channels=1,out_channels=1,kernel_size=2,stride=1,padding=0)
def forward(self,x):
out = self.conv(x)
return out
net = Net()
def rule(epoch):
lamda = math.pow(0.5, int(epoch / 3))
return lamda
optimizer = torch.optim.SGD([{'params': net.parameters(), 'initial_lr': 0.1}], lr = 0.1)
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda = rule)
for i in range(9):
print("lr of epoch", i, "=>", scheduler.get_lr())
optimizer.step()
scheduler.step()
输出如下:
lr of epoch 0 => [0.1]
lr of epoch 1 => [0.1]
lr of epoch 2 => [0.1]
lr of epoch 3 => [0.05]
lr of epoch 4 => [0.05]
lr of epoch 5 => [0.05]
lr of epoch 6 => [0.025]
lr of epoch 7 => [0.025]
lr of epoch 8 => [0.025]
理解LambdaLR的核心在于,自定义规则函数只有一个参数,即当前的epoch,这个参数是scheduler自己传进去的。如果没有特殊指明的话,是从0开始(因为外面的超参last_epoch-1表示-1已结束,因此从0开始),每step一次加1;可以验证如下,在rule里面加一行:
def rule(epoch):
print("current epoch =>", epoch)
lamda = math.pow(0.5, int(epoch / 3))
return lamda
为了方便这里观察,在输出的时候加了换行:
for i in range(9):
print()
print("lr of epoch", i, "=>", scheduler.get_lr())
optimizer.step()
scheduler.step()
结果如下:
current epoch => 0
current epoch => 0
lr of epoch 0 => [0.1]
current epoch => 1
current epoch => 1
lr of epoch 1 => [0.1]
current epoch => 2
current epoch => 2
lr of epoch 2 => [0.1]
current epoch => 3
current epoch => 3
lr of epoch 3 => [0.05]
current epoch => 4
current epoch => 4
lr of epoch 4 => [0.05]
current epoch => 5
current epoch => 5
lr of epoch 5 => [0.05]
current epoch => 6
current epoch => 6
lr of epoch 6 => [0.025]
current epoch => 7
current epoch => 7
lr of epoch 7 => [0.025]
current epoch => 8
current epoch => 8
lr of epoch 8 => [0.025]
current epoch => 9
这个传入的rule函数会在三种情况下调用:
而要恢复训练也很简单,只需要修改两个地方。首先,在恢复训练的情况下,被乘以的初始学习率由必须由优化器的’initial_lr’值指定:
optimizer = torch.optim.SGD([{'params': net.parameters(), 'initial_lr': 0.1}], lr = 0.1)
initial_lr是可以覆盖掉lr的。其次,last_epoch值修改为已完成的epoch数。比如我们想从第5个epoch开始,那么last_epoch就是4,修改如下:
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda = rule, last_epoch=4)
并修改用于训练的while循环:
for i in range(5, 9):
print()
print("lr of epoch", i, "=>", scheduler.get_lr())
optimizer.step()
scheduler.step()
结果如下:
current epoch => 5
current epoch => 5
lr of epoch 5 => [0.05]
current epoch => 6
current epoch => 6
lr of epoch 6 => [0.025]
current epoch => 7
current epoch => 7
lr of epoch 7 => [0.025]
current epoch => 8
current epoch => 8
lr of epoch 8 => [0.025]
current epoch => 9
可以看到同样在第5个epoch结束后调整了学习率。