SWA AveragedModel proper usage

Mario_Parreno: <p>SWA AveragedModel proper usage</p>
<p>Hi, I want to use SWA technique and finally is official at Pytorch! <a href="https://pytorch.org/blog/pytorch-1.6-now-includes-stochastic-weight-averaging/" rel="noopener nofollow ugc">https://pytorch.org/blog/pytorch-1.6-now-includes-stochastic-weight-averaging/</a></p>
<p>But there is something not clear. In the example they post, the averagemodel is created out of the initial training loop. My question is… If we want to average from last model weights, as should be and not initial random weights, we should create the averagemodel just before swa starts, something like this (following example nomenclature):</p>
<pre><code class="lang-auto">if (epoch+1) == swa_start:
    swa_model = AveragedModel(model)
</code></pre>

Mario_Parreno: <p>Up… Not found solution <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=9" title=":confused:" class="emoji" alt=":confused:"> Modifying the page example. I think should be like:</p>
<pre><code class="lang-auto">from torch.optim.swa_utils import AveragedModel, SWALR
from torch.optim.lr_scheduler import CosineAnnealingLR

loader, optimizer, model, loss_fn = ...
swa_model = None
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 &gt; swa_start:
          if swa_model is None:
              swa_model = AveragedModel(model)
          else:
              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)
</code></pre>

jmaronas: <p>Perhaps <a class="mention" href="/u/ptrblck">@ptrblck</a>  can help about this?</p>

ptrblck: <p>I think your code is correct and the initial “checkpoint” would be created after <code>swa_start</code> epochs were already done. Afterwards the update would take place.<br>
I’m unsure, if this is “necessary” or if the posted example would also work fine, since the initial updates would become less important during the training.</p>
<p>Did you run your code and saw a difference using both approaches?<br>
If so, which one performed better?</p>

DDavid: <p>I did like what you did, and it seems you are right.<br>
I checked right after the <code>swa_model = AveragedModel(model)</code>, the weights of the <code>swa_model</code> parameters and <code>model</code> parameters are exactly the same. Doing <code>swa_model = AveragedModel(model)</code> before the loop starts is not a good idea.</p>

hktxt: <p>I wonder how to insert validation phase code?</p>

mikolchon: <p><a class="mention" href="/u/mario_parreno">@Mario_Parreno</a> <a class="mention" href="/u/ptrblck">@ptrblck</a><br>
Do we also need to update the BN stats <em>before</em> validation?<br>
Given that we do need it for testing, I suspect we also need it for validation, correct?</p>

ptrblck: <p>The batchnorm stats will be updated by default during training (the <code>model</code> is by default in training mode or you can additionally call <code>model.train()</code>) while these running stats will be used during validation after calling <code>model.eval()</code>. I’m unsure how you would like to update these stats so could you explain the use case and question a bit more?</p>

mikolchon: <p>Sure! So in SWA two models are maintained: the <code>model</code> and the <code>swa_model</code>. The latter is the averaged model. What we “train” is <code>model</code> (we backprop this, update its weights, etc.), and only every now and then we update <code>swa_model</code> with <code>model</code> by averaging. That’s why the batchnorm stats in <code>swa_model</code> needs separate updating. From the Pytorch website:</p>
<blockquote>
<p>One important detail is the batch normalization. Batch normalization layers compute running statistics of activations during training. Note that the SWA averages of the weights are never used to make predictions during training. So the batch normalization layers do not have the activation statistics computed at the end of training. We can compute these statistics by doing a single forward pass on the train data with the SWA model.</p>
</blockquote>
<p>And here’s the code snippet provided by Pytorch, where we see BN stats of <code>swa_model</code> being updated at the end right before testing:</p>
<pre><code class="lang-auto">for epoch in range(100):
      for input, target in loader:
          optimizer.zero_grad()
          loss_fn(model(input), target).backward()
          optimizer.step()
      if epoch &gt; 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)
</code></pre>
<p>The code snippet doesn’t show the validation step. So my question was whether we should also <code>update_bn()</code> the <code>swa_model</code> for validating it.</p>

ptrblck: <p>Thanks for clarifying the question!<br>
Using the provided code snippet to update the batchnorm stats using the training dataset sounds right. I wouldn’t update the stats using the validation dataset, as I would consider it a data leak (similar to updating them without the SWA util).</p>

mikolchon: <p>Yes, I wouldn’t use the validation dataset. What I meant is to update the BN stats using the training set before running validation, similar to the code example above (<code>loader</code> is the training set).</p>
<p>Ok so looking at <a href="https://github.com/timgaripov/swa/blob/411b2fcad59bec60c6c9eb1eb19ab906540e5ea2/train.py#L143-L161" class="inline-onebox" rel="noopener nofollow ugc">swa/train.py at 411b2fcad59bec60c6c9eb1eb19ab906540e5ea2 · timgaripov/swa · GitHub</a> I think I was right in that we need to update the BN stats of the <code>swa_model</code> before we validate. Relevant lines 158, 159 <a class="mention" href="/u/hktxt">@hktxt</a> <a class="mention" href="/u/ptrblck">@ptrblck</a></p>

chenglu: <p>I was wandering if we update bn with the validation set, will the performance be better? If true, should we also apply the <code>update_bn</code> at the test set?</p>

Fredrik_Opeide: <p>So if we want to compute validation metrics after each training epoch, we need to do an entire second pass through our training data? After every epoch? That can add up to a lot of extra time. Are there any recommended alternative methods?</p>