@izmailovpavel Hi, thanks for the implementation of SWA in PyTorch! Do you find it useful to average the weights during annealing phase? By default anneal_epochs equals to 10 (and cannot be set to 0), so in this snippet for the first 10 epochs after swa_start weights will be averaged during annealing phase.

Hey Daniil, thank you for looking into the implementation! Generally, it's reasonable to average the weights during the annealing phase, but I imagine there could be cases when it's not desirable, e.g. when the learning rate before the annealing is way too high

Thanks for the response! It is just confusing a bit. According to the paper, the purpose of SWA is to average weights during exploration of loss surface (to find high-performing networks), but averaging during annealing would lead to averaging with different (decreasing) learning rates, which paper refers as non suitable (links to experiments in Ruppert, 1988) for improving generalization, due to SGD does not perform very differently under this schedule.

Whether or not averaging in the annealing phase would be desirable would depend on how much the learning rate changes in the annealing phase. Generally, averaging at different learning rates can work well, see e.g. https://arxiv.org/abs/1806.05594. Although it can also be bad, as you said

It may be safer to fix it in the docs as you suggest.

lr_anneal_epochs = 10
swa_scheduler = SWALR(optimizer, swa_lr=0.05, anneal_epochs=lr_anneal_epochs)
...

if i > swa_start + lr_anneal_epochs:
    swa_model.update_parameters(model)

It makes the example a bit more complex... @vincentqb @andrewgordonwilson what do you think?

Is any chance to remove annealing phase from the scheduler? I believe this is the part which accidentally differs from the paper.

It may be safer to fix it in the docs as you suggest.

lr_anneal_epochs = 10
swa_scheduler = SWALR(optimizer, swa_lr=0.05, anneal_epochs=lr_anneal_epochs)
...

if i > swa_start + lr_anneal_epochs:
    swa_model.update_parameters(model)

It makes the example a bit more complex... @vincentqb @andrewgordonwilson what do you think?

What would be an alternate fix that would make this simpler?

I believe the simplest option is to allow SWALR take anneal_epochs=0 by default, thus skip annealing phase. In this case this snippet will be consistent to the paper and can be left as is (but it requires SWALR code modification). Can make a PR if it will work.

Sorry for not following up on this discussion for so long. @Daniil-Osokin I agree that we should allow anneal_epochs=0, and I think it would be great if you could make a pull request for that. I am not sure if we should by default set the number of annealing epochs to zero or not, I am up for discussion. I think both options are reasonable, but I agree that it may be better to set anneal_epochs=0 by default.