Lightning-AI / Lightning-AI/pytorch-lightning

Stochastic weight averaging without the learning rate scheduler

Open
#15,916 1 comment 5 reactions 1 assignee View on GitHub

@justusschock is already working on this.

Since Jan 9, 2023.

callback: swa feature
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.