I have the following code using AdamW
optimizer from Pytorch:
optimizer = AdamW(params=self.model.parameters(), lr=0.00005)
I tried to log in using wandb as follows:
lrs = {f'lr_group_{i}': param_group['lr']
for i, param_group in enumerate(self.optimizer.param_groups)}
wandb.log({"train_loss": avg_train_loss, "val_loss": val_loss, **lrs})
Note that default value for weight_decay
parameter is 0.01
for AdamW
.
When I checked wandb dashboard, it showed the same LR for AdamW even after 200 epochs and it was not decaying at all. I tried this in several runs.
Why the LR decay is not happening?
Also, it was showing LR for only one param group. Why is it so? It seems that I miss something basic here. Can someone point out?
As already hinted at in @xdurch0's comment, I think you have fallen trap to a misconception here. Weight decay and decaying the learning rate are two different things – at least conceptually¹:
Crucially, altering the weight decay value does not alter the learning rate value!
To achieve a decaying learning rate in PyTorch, one usually makes use of a learning rate scheduler, typically in the following way:
step()
on the scheduler (just like calling step()
on the optimizer after processing a training batch). This will trigger the scheduler to calculate the new learning rate value for the next epoch, which it will then use for adjusting the optimizer accordingly.The following is an excerpt of the PyTorch documentation, demonstrating this approach (comments added by me):
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
# 1. Create scheduler, provide it with optimizer
scheduler = ExponentialLR(optimizer, gamma=0.9)
for epoch in range(20):
for input_, target in dataset:
optimizer.zero_grad()
output = model(input_)
loss = loss_fn(output, target)
loss.backward()
optimizer.step()
# 2. Trigger calculation and setting of next learning rate value
scheduler.step()
Here, an ExponentialLR
scheduler is used, which decays the learning rate exponentially, according to a fixed schedule (i.e. only depending on its initial parameters; in particular, the decay factor, but independent of the current network performance). There are also adaptive schedulers, which adjust the learning rate depending on the actual network performance; for example, ReduceLROnPlateau
reduces the learning rate by a given factor as soon as a loss plateau of a given length has been detected.
¹) Mathematically, for some optimizers, learning rate and weight decay are implicitly coupled, which is one of the reasons why AdamW was derived from the Adam optimizer in the first place. For more information, one might want to read the AdamW paper. But this is not the problem described in the question.