Having on_loss_end as callback method


As I am working on something that requires me to modify the loss after it is computed, I noticed that the place it is supposed to be done is the on_backward_begin method. However, this is only called when training, not during validation, which makes it impossible to change the loss when validating. I know it is very niche, but wouldn’t it be better to add a on_loss_end method that is systematically called after loss is computed ? It can be redundant, but it wouldn’t be the only one (on_epoch_begin is too for instance).
For now I patched it locally on my computer and it works fine, so I can suggest a PR if the idea pleases.

EDIT: It raises the question of what preexisting callbacks should be moved here. To me, the only one that is really questionable is Recorder, as one might want to decide whether to record the modified loss or the unmodified one by changing the order. But to me it is fine as it is:

  • If you use on_loss_end, it is because you want the modified loss in validation, so you basically want to record this one. As it will be called before on_backward_begin, it works as intended.
  • If you don’t want to record the changed loss, you probably don’t want to change the validation one, so you can change it in on_backward_begin with an _order over -20. This is current behavior, and it absolutely does not change.
    Therefore I’d change nothing except for allowing to do on_loss_end when it is useful for us.

I don’t really want to change the callbacks system for v1. In v2 the after_loss callback will be called in training and validation.

1 Like

No problem, I’ll stick to my local change for now then. Thanks !