Pytorch-lightning: Add on_after_backwards callback option
馃殌 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
rbracco
All 7 comments
useful!
awaelchli
on 25 Sep 2020
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...
FelixLorenz
on 28 Sep 2020
maybe we should add all the optimizer-related hooks to callbacks too if it makes sense.
What do you think @awaelchli?
rohitgr7
on 1 Oct 2020
@rohitgr7 which hooks do you have in mind?
@FelixLorenz it shouldn't be abstract
awaelchli
on 2 Oct 2020
@awaelchli all of them? on_after_backward, on_before_backward, on_before_zero_grad.
rohitgr7
on 2 Oct 2020
I think that would be nice to have.
awaelchli
on 2 Oct 2020
okay, will add them.
Waiting for more approvals to see if these hooks are good to add or not.
@PyTorchLightning/core-contributors
rohitgr7
on 2 Oct 2020
Related issues
jcreinhold
路
3Comments
polars05
路
3Comments
mmsamiei
路
3Comments
versatran01
路
3Comments
iakremnev
路
3Comments
Most helpful comment
okay, will add them.
Waiting for more approvals to see if these hooks are good to add or not.
@PyTorchLightning/core-contributors