LearningRateMonitor callback causes unexpected changes in step/epoch count with WandBLogger

[rbracco]

🐛 Bug

Using the LRMonitor callback breaks wandb logging by causing the step count to become incorrect. The image below shows varying epoch/step count while overfitting batches with No LR monitor, LearningRateMonitor(logging_interval="epoch") and LearningRateMonitor(logging_interval=None)

Neat-bee-446 does not use the LRMonitor callback and the ratio of step#:epoch# is 1:1
Devout-forest-447 adds as a callback LearningRateMonitor(logging_interval="epoch"), and the ratio of step#:epoch# becomes 2:1
Woven-dew-448 uses the callback LearningRateMonitor() and the ratio of step#:epoch# becomes 3:1

When not overfitting a batch, LearningRateMonitor() has the correct number of steps, but LearningRateMonitor(logging_interval="epoch") and LearningRateMonitor(logging_interval="step") still have double what they should

Also, this doesn't occur with tensorboard, only wandb.

Expected behavior

The logged step count should be correct and not adversely impacted by adding the LRMonitor callback.

Environment

• PyTorch Lightning Version: 1.6.2
• WandB Version: 0.12.5
• PyTorch Version (e.g., 1.11:
• Python version (e.g., 3.8.10):
• OS (e.g., Linux): Linux
• How you installed PyTorch: Pip

Additional context

cc @awaelchli @morganmcg1 @AyushExel @borisdayma @scottire @manangoel99 @rohitgr7

[tanmoyio]

I would like to work on it

[akihironitta]

@rbracco Do you see the issue with other loggers in addition to wandb logger? Would you mind trying the default TensorBoardLogger and seeing if the issue is specific to the logger?

Also, it would be great and very helpful if you could share a script for reproduction.

[manangoel99]

Hi @rbracco ! Engineer from W&B here. The likely reason for this is that pytorch-lightning has its own step counter and wandb has its own step counter as well. The wandb step counter is incremented whenever wandb.log is called which can be multiple times within the same lightning step.

For example:

def training_step(self, batch):
    self.log("a", 1)
    self.log("b", 2)

This is 1 training step but this would increment the wandb step twice.
Hence, adding the LRMonitor callback would add another log call which would further increment the wandb step.

[awaelchli]

Hey @rbracco if what @manangoel99 is the case, you can change the x-axis to "trainer step" in the UI (top right) either globally or per plot.

[manangoel99]

@rbracco Please let me know if the suggested solution works!

[rbracco]

Thank you @awaelchli and @manangoel99 that worked great. Here is the updated chart once I change x-axis to trainer step. Closing!

[awaelchli]

yay :) Happy logging!