Validation code runs before `training_epoch_end()` in v1.4 and breaks SWA

[chanshing]

This PR #7357 which was recently introduced in v1.4 breaks cases where one needs to somehow update the model before running validation/testing, such as with the increasingly popular SWA (stochastic weight averaging) method.

From https://pytorch.org/blog/pytorch-1.6-now-includes-stochastic-weight-averaging/

from torch.optim.swa_utils import AveragedModel, SWALR
from torch.optim.lr_scheduler import CosineAnnealingLR

loader, optimizer, model, loss_fn = ...
swa_model = AveragedModel(model)
scheduler = CosineAnnealingLR(optimizer, T_max=100)
swa_start = 5
swa_scheduler = SWALR(optimizer, swa_lr=0.05)

for epoch in range(100):
      for input, target in loader:
          optimizer.zero_grad()
          loss_fn(model(input), target).backward()
          optimizer.step()
      if epoch > swa_start:
          swa_model.update_parameters(model)
          swa_scheduler.step()
      else:
          scheduler.step()

# Update bn statistics for the swa_model at the end
torch.optim.swa_utils.update_bn(loader, swa_model)
# Use swa_model to make predictions on test data 
preds = swa_model(test_input)

The code snippet above lacks a validation step but shows a test step. Basically, the swa_model bn statistics have to be updated by running one forward loop on the dataloader before testing/validating.

To employ SWA in my lightning code, I have something as follows:

def training_epoch_end(self, training_step_outputs=None):
    if self.current_epoch >= self.swa_start:
        self.swa_model.update_parameters(self.model)
        torch.optim.swa_utils.update_bn(loader, self.swa_model)  # update bn stats

This used to run before the validation step, but in v1.4 it now runs after the validation step. In practice, this means that the validation metrics are delayed by one epoch.

I'm aware there's work on incorporating SWA into Lightning https://pytorch-lightning.readthedocs.io/en/latest/extensions/generated/pytorch_lightning.callbacks.StochasticWeightAveraging.html but it's currently in beta and lacks documentation.

In general, is there now (v1.4) any way to insert code to run after a training epoch end and before the validation step?

[carmocca]

any way to insert code to run after a training epoch end and before the validation step?

You can check if updating the parameters is necessary in the on_validation_start or on_validation_epoch_start and then do it.

See the pseudocode with all hooks here:

https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#hooks

but it's currently in beta and lacks documentation.

Curious, what kind of documentation are you missing?
Our docs about it are short since we just reference PyTorch's.

You can also take a look at our implementation to take ideas if you need to

https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/callbacks/stochastic_weight_avg.py

[tchaton]

Dear @chanshing,

Definitely curious to hear what you think is missing with PyTorch Lightning built-in callback :)

I would advice you to have a look at how it has been implemented: https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/callbacks/stochastic_weight_avg.py#L34

Best,
T.C

[chanshing]

You can check if updating the parameters is necessary in the on_validation_start or on_validation_epoch_start and then do it.

See the pseudocode with all hooks here:

https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#hooks

Thanks for the suggestions. Indeed a solution can be to add the model update code in on_validation_start but it doesn't seem right for on_validation_start to have a side effect. There's also some code duplication as the same code has to be added in training_epoch_end or on_train_end at the latest.

Another option is to put it in on_train_batch_end but the issue is that the model update is very expensive (full loop through the dataloader to recalibrate the BN stats). I could include an if so that it runs at some interval, but this should be in sync with the should_check_val condition. Somehow it seems there should be some hook separating training and validation, although I understand that it might not be worth changing things just for SWA.

Definitely curious to hear what you think is missing with PyTorch Lightning built-in callback :)

I would advice you to have a look at how it has been implemented: https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/callbacks/stochastic_weight_avg.py#L34

I tried to follow https://pytorch-lightning.readthedocs.io/en/latest/extensions/generated/pytorch_lightning.callbacks.StochasticWeightAveraging.html and I think it could help to have an example. There it says that Trainer(stochastic_weight_avg=True) is all I need but I don't know where and how to set the other params (swa_epoch_start, annealing_epochs, etc.). Do I have to pass a callback and use stochastic_weight_avg=True?

[carmocca]

There it says that Trainer(stochastic_weight_avg=True) is all I need but I don't know where and how to set the other params (swa_epoch_start, annealing_epochs, etc.).

Opening a PR to improve it!

Do I have to pass a callback and use stochastic_weight_avg=True?

No, use one or the other

[stale]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!