Lightning-AI / Lightning-AI/pytorch-lightning
Stochastic weight averaging without the learning rate scheduler
@justusschock is already working on this.
Since Jan 9, 2023.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
## 🚀 Feature
The ability to use stochastic weight averaging _without_ the SWALR learning rate scheduler.
### Motivation
Unless I'm mistaken, the current implementation of `StochasticWeightAveraging` will do the averaging, anneal the learning rate to a constant value, and update the batch norms. This is wonderful, but what if I want to experiment with a different learning rate schedule while SWA is active (as they do in the paper, using Cyclic learning rates for some models).
### Pitch
It would be nice to have a parameter on the `StochasticWeightAveraging` callback, `anneal_to_constant_lr` that defaults to `True` (resulting in the current behaviour) but that I can set to `False` so that whatever learning rate schedules I have in place are not affected by the use of this callback (but the averaging, batch norm, and switching which model is used for the validation/test step all stay as is).
### Alternatives
I'm new to Lightning, but assume I can implement the underlying PyTorch SWA components to get this effect, I just haven't worked out how yet and was hoping I could use this callback.
cc @borda @carmocca
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.