Allow logging in `on_train_end`

[lminer]

🚀 Feature

I would like to be able to log the metrics of the best performing epoch at the end of training. This is especially useful for hyperparameter sweeps. on_train_end seems like a natural place for this.

[Atharva-Phatak]

@lminer You can accumulate all the metrics from train_step and just log them.
Are you referring to any specific logging service like WANDB, etc or just in general self.log method which lightning provides ?

[lminer]

I want to have some way to run the general self.log method at the very end of training in order to sum up the run. on_train_end seems to be the place for this, but if there's somewhere better, it would be great to know.

[awaelchli]

@carmocca I don't remember the exact criterions for when we allow to log in a hook or not. Perhaps there is a good reason, probably related to reduction as on_epoch/on_step is not meaningful post the epoch loops, but I'm not sure.

@lminer You should be able to work around this issue for now by calling the logger directly (this is always possible):

self.logger.experiment.add_scalar("best", value)  # e.g. for tensorboard logger

[carmocca]

That's correct @awaelchli. on_train_end is the last hook to run and all the internal machinery has been torn down at this point. self.logging in it wouldn't work.

[awaelchli]

Yes, makes sense. Then I suppose #15143 (comment) is becoming more like our official recommendation for such cases.

Do you think it is meaningful to add a hint to our error messages that say "it is not allowed to self.log in x", where we would recommend "you can alternatively log with logger.experiment..."?

[carmocca]

Sure

[Al3xDo]

Hi @awaelchli, I would like to work on this issue.