How to see the adapted learning rate for Adam in pytorch?

Viewed 2273

There are many different optimizers with adaptive learning rate methods. Is it possible to see the adapted value of the initial learning rate for Adam?

Here is a similar question about Adadelta and the answer was to search for ["acc_delta"] key, but Adam has no that key.

2 Answers

AFAIK there is no super easy way to do this. However, you can recalculate the current learning rate of a certain paramter using the implementation of Adam in PyTorch: https://pytorch.org/docs/stable/_modules/torch/optim/adam.html

I came up with this minimal working example:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable

def get_current_lr(optimizer, group_idx, parameter_idx):
    # Adam has different learning rates for each paramter. So we need to pick the
    # group and paramter first.
    group = optimizer.param_groups[group_idx]
    p = group['params'][parameter_idx]

    beta1, _ = group['betas']
    state = optimizer.state[p]

    bias_correction1 = 1 - beta1 ** state['step']
    current_lr = group['lr'] / bias_correction1
    return current_lr

x = Variable(torch.randn(100, 1)) #Just create a random tensor as input
model = nn.Linear(1, 1)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
niter = 20
for _ in range(0, niter):
    out = model(x)

    optimizer.zero_grad()
    loss = criterion(out, x) #Here we learn the identity mapping
    loss.backward()
    optimizer.step()
    group_idx, param_idx = 0, 0
    current_lr = get_current_lr(optimizer, group_idx, param_idx)
    print('Current learning rate (g:%d, p:%d): %.4f | Loss: %.4f'%(group_idx, param_idx, current_lr, loss.item()))

which should output something like this:

Current learning rate (g:0, p:0): 0.0100 | Loss: 0.5181
Current learning rate (g:0, p:0): 0.0053 | Loss: 0.5161
Current learning rate (g:0, p:0): 0.0037 | Loss: 0.5141
Current learning rate (g:0, p:0): 0.0029 | Loss: 0.5121
Current learning rate (g:0, p:0): 0.0024 | Loss: 0.5102
Current learning rate (g:0, p:0): 0.0021 | Loss: 0.5082
Current learning rate (g:0, p:0): 0.0019 | Loss: 0.5062
Current learning rate (g:0, p:0): 0.0018 | Loss: 0.5042
Current learning rate (g:0, p:0): 0.0016 | Loss: 0.5023
Current learning rate (g:0, p:0): 0.0015 | Loss: 0.5003
Current learning rate (g:0, p:0): 0.0015 | Loss: 0.4984
Current learning rate (g:0, p:0): 0.0014 | Loss: 0.4964
Current learning rate (g:0, p:0): 0.0013 | Loss: 0.4945
Current learning rate (g:0, p:0): 0.0013 | Loss: 0.4925
Current learning rate (g:0, p:0): 0.0013 | Loss: 0.4906
Current learning rate (g:0, p:0): 0.0012 | Loss: 0.4887
Current learning rate (g:0, p:0): 0.0012 | Loss: 0.4868
Current learning rate (g:0, p:0): 0.0012 | Loss: 0.4848
Current learning rate (g:0, p:0): 0.0012 | Loss: 0.4829
Current learning rate (g:0, p:0): 0.0011 | Loss: 0.4810

Note that monitoring the learning rate of every individual paramter is probably not feasible nor helpful for larger models.

To build upon the answer by Elias E., I would consider not only the bias correction part of the effective learning rate, but also the per parameter normalization depending on the second momentum, so an update to the original function would be:

def get_current_lr(optimizer, group_idx, parameter_idx):
    # Adam has different learning rates for each paramter. So we need to pick the
    # group and paramter first.
    group = optimizer.param_groups[group_idx]
    p = group['params'][parameter_idx]

    beta1, _ = group['betas']
    state = optimizer.state[p]

    bias_correction1 = 1 - beta1 ** state['step']
    current_lr = group['lr'] / bias_correction1 / torch.sqrt(state['exp_avg_sq'] + 1e-8)
    return current_lr

where you devide by the exponential moving average of the squared gradient.

Related