Ability to specify step for logged metrics

[festeh]

🚀 Feature

Add option to specify step for logged metrics

Motivation

After calculating some metric on n-th epoch I'd like to put corresponding mark (n) on x-axis. As I see in the source code, there's no obvious way to do this. Instead something like n * num_batches will be shown, like in the figure where the loss is calculated after each epoch.

Pitch

I'd love to see 0,1,2,... ticks on x-axis.
My proposal is to allow user to add step key in log dict, if this key is presented, corresponding value would be used in x-axis. I can make a PR if you agree with this idea.

Alternatives

It would be interesting to know how this could be achieved in other way.

[awaelchli]

An alternative would be to call the specific experiment object from the logger, e.g., for TensorBoard logger:

self.logger.experiment.add_scalar("myloss", loss, global_step=...)

I think that should work. It will still log the other metrics you return in the dict with the regular step values.

[festeh]

You're right, this is a vaild way, when I opened this issue I did not thought about it. Still I think that this approach is a bit complex, as for such a simple task you need to cast your metrics to float and possibly write a loop with add_scalar instead of just passing metrics in log dict.