Avoid zero_grad()

[cowwoc]

https://ai.plainenglish.io/best-performance-tuning-practices-for-pytorch-3ef06329d5fe section Zeroing The Gradients Efficiently proposes an alternative that they claim is more efficient than invoking zero_grad(). Has this been tried before?

[ananthsub]

There's a set_to_none flag as an option on the optimizer's zero_grad: https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html https://pytorch.org/docs/stable/_modules/torch/optim/optimizer.html#Optimizer.zero_grad

Have you tried using this?

[cowwoc]

@ananthsub No, I have not tried this myself but I assume it works. All my projects build on top of Lightning which does not provide a mechanism to set this flag, so I can't play around with this.

What changes would you need to make in Lightning? Is it as simple as adding set_to_none=True in the few places that invoke optimizer.zero_grad() or would further changes be needed?

[rohitgr7]

all you need to do is update one of the lightning hooks
https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#optimizer-zero-grad

[cowwoc]

Oh. That's amazing! I wasn't aware of that hook. Thank you!