PyTorch - 如何在训练期间获得学习率？

Posted 2023-03-12

【中文标题】PyTorch - 如何在训练期间获得学习率？

【英文标题】：PyTorch - How to get learning rate during training?

【发布时间】：2019-03-10 16:30:31

【问题描述】：

在训练时，我想知道 learning_rate 的值。 我该怎么办？

这是我的代码，如下所示：

my_optimizer = torch.optim.SGD(my_model.parameters(), 
                               lr=0.001, 
                               momentum=0.99, 
                               weight_decay=2e-3)

谢谢。

【参考方案1】：

对于您给出的示例中只有一个参数组，您可以使用此函数并在训练期间调用它以获得当前的学习率：

def get_lr(optimizer):
    for param_group in optimizer.param_groups:
        return param_group['lr']

【讨论】：

在没有动量和调度器的情况下，在 SGD 上将 lr 设置为 0.1。在训练期间得到了 7.2 的东西。我认为显示的值不正确

【参考方案2】：

或者，您可以将lr_scheduler 与优化器一起使用，并简单地调用内置的lr_scheduler.get_lr() 方法。

这是一个例子：

my_optimizer = torch.optim.Adam( my_model.parameters(), 
                                 lr = 0.001, 
                                 weight_decay = 0.002)

my_lr_scheduler = torch.optim.lr_scheduler.StepLR( my_optimizer, 
                                                step_size = 50, 
                                                gamma = 0.1)

# train
...
my_optimizer.step()
my_lr_scheduler.step()

# get learning rate
my_lr = my_lr_scheduler.get_lr()
# or
my_lr = my_lr_scheduler.optimizer.param_groups[0]['lr']

使用lr_scheduler 的额外好处是可以更好地控制随时间变化的 lr； lr_decay 等 lr_scheduler args 参考 pytorch docs。