Pytorch lr_scheduler.LambdaLR()的简单理解与用法

官方文档：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)

参数：

• optimizer：被调整学习率的优化器
• lr_lambda：用户自定义的学习率调整规则。可以是lambda表达式，也可以是函数
• last_epoch：当前优化器的已迭代次数，后文我们将其称为epoch计数器。默认是-1，字面意思是第-1个epoch已完成，也就是当前epoch从0算起，从头开始训练。如果是加载checkpoint继续训练，那么这里要传入对应的已迭代次数
• verbose：是否在更新学习率时在控制台输出提醒

一个例子如下。考虑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函数会在三种情况下调用：

• scheduler被创建时。因此可以看到一开始输出了个current epoch => 0。
• 调用get_lr()方法时。这里有一点特别需要注意，get_lr()本质上就是将内置的epoch计数器传进rule方法，然后拿一个lambda出来，lambda乘以初始学习率就是当前学习率。而get_last_lr()则不会有这一过程。
• 调用scheduler.step()方法时。这个时候内置的epoch计数器会加一，并依据规则更新学习率。

而要恢复训练也很简单，只需要修改两个地方。首先，在恢复训练的情况下，被乘以的初始学习率由必须由优化器的’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结束后调整了学习率。