How to sync learning rate of lr_scheduler with that of optimizer

[Infin1te2021]

Hey, I'm trying to use CosineLRScheduler, but I noticed the optimizer learning rate isn't in sync with lr_scheduler. The optimizer learning rate I'm getting is warmup_lr_init. The learning rate calculated by lr_scheduler seems correct. Here's my sample for checking the learning rate. And I follow the same way to train my model.

timm = 1.0.15
torch = 2.5.1+cu124

import torch
from torch import nn
from timm.scheduler.cosine_lr import CosineLRScheduler
import matplotlib.pyplot as plt

epoch = 100
warmup_epoch = 5
itr = 125
num_steps = itr * epoch
warmup_step = itr * warmup_epoch
warmup_prefix = True
t_initial=(num_steps-warmup_step) if warmup_prefix else num_steps

# Example
class Example(nn.Module):
  def __init__(self, num_class):
    super().__init__()
    self.proj = nn.Linear(10, 20)
    self.proj1 = nn.Linear(20, num_class)

  def forward(self, x):
    N, C= x.size()
    x = self.proj(x)
    x = self.proj1(x)
    return x

batch_size = 32
data = torch.randn(batch_size, 10) 
model = Example(10)

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

class DebugCosineLRScheduler(CosineLRScheduler):
  def _get_lr(self, t: int):
    lrs = super()._get_lr(t)
    return lrs

lr_scheduler = DebugCosineLRScheduler(optimizer, 
                                t_initial=t_initial,
                                lr_min = 1e-5,
                                warmup_lr_init = 1e-7,
                                warmup_t=warmup_step,
                                warmup_prefix= False,
                                cycle_limit=1,
                                t_in_epochs=False)
clrs = []
lrs = []
for i in range(num_steps):
  optimizer.zero_grad()
  output = model(data)
  loss = output.mean()
  loss.backward()
  optimizer.step()
  lr_scheduler.step(i)

  calculated_lr = lr_scheduler._get_lr(i)[0]  # get the learning rate in each step from lr_scheduler
  clrs.append(calculated_lr)
  current_lr = optimizer.param_groups[0]['lr'] # get the learning rate in each step from optimizer
  lrs.append(current_lr)

  # assert abs(current_lr - calculated_lr) < 1e-9, "Learning rate not equal"


plt.plot(clrs, label="Calculated LR", color="blue")
plt.plot(lrs, label="Optimizer LR", color="orange")
plt.xlabel("Step")
plt.ylabel("Learning Rate")
plt.title("CosineLRScheduler Learning Rate")
plt.legend()
plt.show()

[Infin1te2021]

I should call lr_scheduler.step_update() instead of lr_scheduler.step()

def step(self, epoch: int, metric: float = None) -> None:
  self.metric = metric
  values = self._get_values(epoch, on_epoch=True)
  if values is not None:
    values = self._add_noise(values, epoch)
    self.update_groups(values)

def step_update(self, num_updates: int, metric: float = None):
  self.metric = metric
  values = self._get_values(num_updates, on_epoch=False)
  if values is not None:
    values = self._add_noise(values, num_updates)
    self.update_groups(values)

Now the plot is correct.