Changing the computed gradients before backprop

[sebastiangonsal]

I want to add noise to the gradients in pytorch lightning . Specifically, something similar to this paper: https://arxiv.org/pdf/1511.06807.pdf . Basically, I would compute the gradients and before the call to backward, i want to add noise.

What is the best way to achieve this in pytorch lightning?

[ananthsub]

Does the on_before_backward hook work for your use case? https://github.com/PyTorchLightning/pytorch-lightning/blob/24f0124ddd478e18b76d64a79f1d7b341a0f8646/pytorch_lightning/core/hooks.py#L298-L305

You can override this in your LightningModule

If you need absolute control over the optimization process (forward, backward, optimizer steps), you can use manual optimization: https://pytorch-lightning.readthedocs.io/en/latest/common/optimizers.html#manual-optimization

[tchaton]

Dear @sebastiangonsal,

From your paper, it seems you might want to call backward which computes the gradients and then add some noise right ?

Therefore, you could override the before_optimizer_step and add noise to all the params.grad.

class TestModel(LightningModule):

    def on_before_optimizer_step(self):
        for param in self.parameters():
            param.grad += torch.rand_like(param.grad)

If you want this re-usable, you could just move this to a callback instead.