Add on_after_backwards callback option

[rbracco]

🚀 Feature

Add an on_after_backward callback option.

Motivation

Currently there are no callbacks for after the backwards step so non-essential code like logging gradients clutters the LightningModule.

Pitch

Expand the callback options to include a hook for on_after_backward() to execute callbacks immediately after the backwards pass

Alternatives

Sticking my non-essential gradient logging code in my LightningModule

[awaelchli]

useful!

[FelixLorenz]

We nearly need the same feature: on_before_backward as an abstract method in the Callback base class would be awesome!

Currently we are using the same alternative (in the training_step) but this couples one last bit of our custom callback with the code of the lightning module...

[rohitgr7]

maybe we should add all the optimizer-related hooks to callbacks too if it makes sense.
What do you think @awaelchli?

[awaelchli]

@rohitgr7 which hooks do you have in mind?
@FelixLorenz it shouldn't be abstract

[rohitgr7]

@awaelchli all of them? on_after_backward, on_before_backward, on_before_zero_grad.

[awaelchli]

I think that would be nice to have.

[rohitgr7]

okay, will add them.
Waiting for more approvals to see if these hooks are good to add or not.
@PyTorchLightning/core-contributors