SWA models are not working

thorstenwagner: <p>Hi,</p>
<p>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.</p>
<p>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.</p>
<p>I train a mixed-precision model using DDP. I think, the way I calculate the model is actually straightforward:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/MPI-Dortmund/tomotwin-cryoet/blob/stat-weight-averaging/tomotwin/modules/training/torchtrainer.py#L381">
  <header class="source">

      <a href="https://github.com/MPI-Dortmund/tomotwin-cryoet/blob/stat-weight-averaging/tomotwin/modules/training/torchtrainer.py#L381" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/MPI-Dortmund/tomotwin-cryoet/blob/stat-weight-averaging/tomotwin/modules/training/torchtrainer.py#L381" target="_blank" rel="noopener nofollow ugc">MPI-Dortmund/tomotwin-cryoet/blob/stat-weight-averaging/tomotwin/modules/training/torchtrainer.py#L381</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="371" style="counter-reset: li-counter 370 ;">
          <li>        scaler.step(self.optimizer)</li>
          <li>        scaler.update()</li>
          <li>        desc_t = f"Training (loss: {np.mean(running_loss[-20:]):.4f}) "</li>
          <li></li>
          <li>        t.set_description(desc=desc_t)</li>
          <li></li>
          <li>    training_loss = np.mean(running_loss)</li>
          <li>    self.last_loss = training_loss</li>
          <li>    return training_loss</li>
          <li></li>
          <li class="selected">def train(self) -&gt; nn.Module:</li>
          <li>    """</li>
          <li>    Trains the model and returns it.</li>
          <li>    :return: Trained model</li>
          <li>    """</li>
          <li>    if self.training_data is None:</li>
          <li>        raise RuntimeError("Training data is not set")</li>
          <li></li>
          <li>    train_loader, test_loader = self.get_train_test_dataloader()</li>
          <li>    from torch.optim.swa_utils import AveragedModel, get_ema_multi_avg_fn</li>
          <li>    self.swa_model = AveragedModel(self.model, multi_avg_fn=get_ema_multi_avg_fn(0.999))</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>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:</p>
<p>not averaged model: ‘conv_layer0.0.weight’</p>
<p>averaged model: ‘module.conv_layer0.0.weight’</p>
<p>Could that be already the problem?</p>
<p>Happy to get any new insights <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p>Best,<br>
Thorsten</p>

thorstenwagner: <p>I think I can narrow down the problem a bit.</p>
<p>The training runs with the old “DataParallel” mode.</p>
<p>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.</p>
<p>Any idea what might be causing this?</p>
<p><a class="mention" href="/u/ptrblck">@ptrblck</a> maybe?</p>
<p>Best,<br>
Thorsten</p>
<p>PS: The topic now might be better placed into the “Distributed” category.</p>

thorstenwagner: <p>Solved. It was just the module. prefix. Removed it and now it works <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>