Add training step callback hook before .backward

[GR4HAM]

🚀 Feature

The proposal is to add a callback hook for the training step before .backward. Current hooks run either after .backward (on_after_backward) or after optimizer.step (on_train_batch_end, on_batch_end). The proposed callback hook would complement the existing hooks.

Motivation

Adding this hook allows one to create callbacks related to loss calculations as well as batch operations for loss calculations. The current hooks cannot be used for this because the automatic pl optimizer is already called before the existing hooks.

Adding the proposed hook would increase the flexibility and usability of pytorch lightning without modifying existing functionality.

Pitch

Concretely, it would suffice to add an additional callback hook, named for example "on_before_backward", that takes in trainer and pl_module as arguments.

Alternatives

Without the proposed callback, the alternative for the end user is to modify their own code to create their own hook in all the pl modules they use.

[tchaton]

Dear @GR4HAM,

Do you want to give it a try ?

Here are the steps:

• 1. Add the new hook to the pytorch_lightning/callbacks/base.py::Callback
• 2. Add the new hook inside pytorch_lightning/trainer/callbacks.py
• 3. Find in the TrainLoop where the hook should be call and use self.trainer.call_hook("on_before_backward")
• 4. Add the new callback to the FXValidator (same as training_step)
• 5: Resolve failing tests. Lambda Callback + FXValidator.

Best,
T.C