SWA models are not working

thorstenwagner · April 29, 2024, 4:28pm

Hi,

I’m trying to get SWA (EMA) working for my model. But whatever I try, the model I get afterwards does not work. Well, it technically works but the predictions are simply wrong.

The way I currently do it is, that I start from my current best model as checkpoint and train for another 20 epochs with a reduced learning rate. And then I use the averaged model.

I train a mixed-precision model using DDP. I think, the way I calculate the model is actually straightforward:

github.com

MPI-Dortmund/tomotwin-cryoet/blob/stat-weight-averaging/tomotwin/modules/training/torchtrainer.py#L381


      
                  scaler.step(self.optimizer)
                  scaler.update()
                  desc_t = f"Training (loss: {np.mean(running_loss[-20:]):.4f}) "
          
                  t.set_description(desc=desc_t)
          
              training_loss = np.mean(running_loss)
              self.last_loss = training_loss
              return training_loss
          
          def train(self) -> nn.Module:
              """
              Trains the model and returns it.
              :return: Trained model
              """
              if self.training_data is None:
                  raise RuntimeError("Training data is not set")
          
              train_loader, test_loader = self.get_train_test_dataloader()
              from torch.optim.swa_utils import AveragedModel, get_ema_multi_avg_fn
              self.swa_model = AveragedModel(self.model, multi_avg_fn=get_ema_multi_avg_fn(0.999))

I was wondering if the way I serialize the model could be a problem, because the layer keys in the model_state_dict is a bit different than the layer names in the not averaged model:

not averaged model: ‘conv_layer0.0.weight’

averaged model: ‘module.conv_layer0.0.weight’

Could that be already the problem?

Happy to get any new insights

Best,
Thorsten

thorstenwagner · May 3, 2024, 6:43am

I think I can narrow down the problem a bit.

The training runs with the old “DataParallel” mode.

However, the “prediction” runs with “DistributedDataParallel”. So far this has not been a problem, but it turns out that running the averaged model with DDP gives me corrupted results, while running the same model with DP gives me decent results.

Any idea what might be causing this?

@ptrblck maybe?

Best,
Thorsten

PS: The topic now might be better placed into the “Distributed” category.

thorstenwagner · May 3, 2024, 11:37am

Solved. It was just the module. prefix. Removed it and now it works