Support checkpoint save and load with Stochastic Weight Averaging

[adamreeve]

What does this PR do?

Partly addresses #6074 by supporting saving and loading the StochasticWeightAveraging callback data in checkpoints. Support for using SWA during validation will be done as a follow up PR.

Does your PR introduce any breaking changes? If yes, please list them.

No

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG**? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

• Make sure you had fun coding 🙃

[rohitgr7]

Perfect!
Great work!

Apologies for the delayed review! but it was complex and you pulled it out smartly 😃 🚀

[awaelchli]

@rohitgr7 Will this go into 1.7.x?

[rohitgr7]

@awaelchli yes

[kampelmuehler]

@adamreeve
Am I correct that the "follow up PR" on using the _average_model for validation hasn't yet been done?

For those looking to do it manually, this would be a very rough sketch

    def on_fit_start(self):
        for cb in self.trainer.callbacks:
            if isinstance(cb, StochasticWeightAveraging):
                self.swa = cb

    def validation_step(self, batch, batch_idx):
        x, y = batch

        loss = self.loss_fn(self(x), y)
        print(f"{loss.item():.02f}", end="")
        if self.swa._initialized:
            loss_swa = self.loss_fn(self.swa._average_model(x), y)
            print(f" {loss_swa.item():.02f}", end="")
        print()
        return loss

[kampelmuehler]

@adamreeve
I think that these features should also somehow be documented. Currently the docs include nothing on how to actually use the SWA model (or that indeed only in the end of training the weights are transferred to the original model). What do you think?

Edit:
Couple of users unsure about SWA usage
https://lightning.ai/forums/t/how-to-implement-swa/1761
https://lightning.ai/forums/t/stochasticweightaveraging-validation-logging-and-checkpoints/2023/2

[kampelmuehler]

Also:
The checkpointing is not in any sense automated, or is it?
Only way I see thus far is to manually load the checkpoint, get the state_dict for swa and load it manually into the callback instance.

If I just declare the model, add an swa callback to the trainer and pass the checkpoint path to trainer.fit() it won't load the swa callback states.

EDIT:
I was wrong, passing to trainer.fit() correctly loads the callback state as well, trainer.predict(), however, will not load the swa callback state. Which might be intentional?

[kampelmuehler]

cc @awaelchli

[adamreeve]

Hi @kampelmuehler

Am I correct that the "follow up PR" on using the _average_model for validation hasn't yet been done?

Yes, that was originally part of this PR but the scope was reduced, so #6074 probably shouldn't have been closed. I had started working towards that in a separate branch (https://github.com/adamreeve/pytorch-lightning/commits/swa_validation) but that's now quite far behind master and doesn't include all the changes from this PR. Supporting batch normalization in conjunction with SWA made that a lot more complicated, but I think it was working. I don't currently have any plans to continue with that work.

This PR didn't really add any new user-visible feature but just fixed checkpointing to work correctly with SWA so that training could be resumed. Eg. it fixed #11665. But I agree that the documentation could be improved to better explain how SWA works.

It's a while since I looked at this now so I'm not sure whether trainer.predict() not loading the SWA parameters is intentional or just a limitation of the current approach, but it sounds consistent with the behaviour of the averaged parameters not being transferred until training is completed.

[kampelmuehler]

Hi @adamreeve - thanks for the quick response and all the insights!