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 @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?