How to implement manual gradient computation and scaling in Pytorch Lightning?

[FeryET]

Hi.

I'm trying to compute gradient penalty in my training loop, and from what I've read in Pytorch's documentation, gradient penalty should be computed in a context in which gradients are scaled and no kind of mixed precision acceleration is going on. How can I achieve this in pytorch lightning? From what I've read manual_backward is still wrapped via mixed precision accelerations.

Thanks!

[rohitgr7]

hey @FeryET !!

I am not familiar with gradient penalty. Can you share some references to the PyTorch doc and gradient penalty? I can check and help you here.

[FeryET]

https://pytorch.org/docs/stable/notes/amp_examples.html#gradient-penalty

This is an example of using gradient penalty in automatic mixed precision. As you can see the loss is scaled in order to compute the gradients and then the gp gets rescaled to match the half precision loss.