Make logging_mode (on_step or on_epoch) available to loggers

[SiddhantSadangi]

Description & Motivation

Lightning lets you log on_step, or on_epoch. However, even when using on_epoch, the global step is logged on the x-axis, not the epoch number.

Loggers can manually extract epoch from the logged metrics and pass that as the step. However, this would break on_step logging.

Making the logging mode for a logging call available would let loggers conditionally use either the global_step or current_epoch as the step for that particular metric.

Pitch

No response

Alternatives

No response

Additional context

No response

cc @lantiga @Borda

[lantiga]

Hey @SiddhantSadangi, this is very reasonable. Would you like to send a PR?
The key thing to ensure is that we don't break the behavior of existing loggers.

[SiddhantSadangi]

Hey @SiddhantSadangi, this is very reasonable. Would you like to send a PR?

I can start working on it once I've some bandwidth. It can take a couple of months, though, due to present commitments :(

[lantiga]

That would be great @SiddhantSadangi! Of course, let's circle back in a couple of weeks and see where you are in terms of bandwidth.

BTW anyone interested in helping?

[SiddhantSadangi]

I'll share this thread with a few Neptune users who had initially requested this. Meanwhile, can we add the Help wanted label here to invite contributions?

[DanTremonti]

In my case, the mentioned change in x-axis measure of the data points started appearing while upgrading lightning version from 2.0 to 2.2. Since this is making logs not comparable, I'm looking for

1. the release version that contains this change
2. the motivation behind bringing about this change
3. any workaround that can be applied on my end to temporarily revert this behaviour