Adapting SWA for pytorch lightning.

Hello! I’m currently trying to add SWA support to my lightning module, following this post https://pytorch.org/blog/stochastic-weight-averaging-in-pytorch/ And I have problems with loading averaged weights after I finish training. I do it like this:

def on_train_end(self): #hook for SWA weights swapping
        if self.swa_enabled:
            print("Swapping model weights to SWA averages.")
            self.trainer.optimizers[0].swap_swa_sgd()
            print("Saving model with SWA weights to another checkpoint")
            self.save_checkpoint(os.path.join(self.trainer.checkpoint_callback.dirpath,f"swa_model_epoch={self.current_checkpoint}.ckpt"))

And bump into following problems:

1. It returns error message 'MyCoolModel' object has no attribute 'save_checkpoint' on the line self.save_checkpoint....., instead of saving the model as stated in https://pytorch-lightning.readthedocs.io/en/latest/weights_loading.html 2)Even if I succeed with the 1st issue, I’m still required to update BN instances. As I believe, the best way to do this is to just run forward pass over the whole training data. So, my question is: is there a simple way to run 1 epoch though training data without writing my own loops? Thanks in advance

[jeremyjordan]

It returns error message ‘MyCoolModel’ object has no attribute ‘save_checkpoint’ on the line self.save_checkpoint…

i believe the documentation is wrong, save_checkpoint is a method for the Trainer class and not the LightningModule class.

https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/trainer/training_io.py#L247

[jeremyjordan]

@RafailFridman feel free to reopen this issue if you’re still having problems