SyncBatchNorm with SWA

James_W: <p>Hi thre,</p>
<p>I was wondering if there was any docs on how to use SyncBatchNorm with SWA. I have a mobilenet pretrained model which I converted into SyncBatchnorm using:</p>
<pre><code class="lang-auto">    model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
</code></pre>
<p>and then do the DDP stuff and then I tried to update batch stats at the end of the training using the uility fiunction like so:</p>
<pre><code class="lang-auto">                    gpu_device=torch.device('cuda:1')
                    swa_model=swa_model.to(gpu_device)
                    print('==&gt;  UPDATING BATCH STATS')
                    torch.optim.swa_utils.update_bn(train_loader, swa_model,gpu_device)
                    print('==&gt;  FINISHED UPDATING BATCH STATS')
</code></pre>
<p>It starts updating the batch stats and then nothing happens… literally nothing… no error, no exit… it just sits there…</p>
<p>When I remove syncbatchnorm and train and update stats it all works perfectly.</p>
<p>I browsed through the source of the above:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/34bdd46e6e6a15dfa1c765c1fb8d9d7bc02dc102/torch/optim/swa_utils.py#L136">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/34bdd46e6e6a15dfa1c765c1fb8d9d7bc02dc102/torch/optim/swa_utils.py#L136" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/34bdd46e6e6a15dfa1c765c1fb8d9d7bc02dc102/torch/optim/swa_utils.py#L136" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/34bdd46e6e6a15dfa1c765c1fb8d9d7bc02dc102/torch/optim/swa_utils.py#L136</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="126" style="counter-reset: li-counter 125 ;">
          <li>            device = p_swa.device</li>
          <li>            p_model_ = p_model.detach().to(device)</li>
          <li>            if self.n_averaged == 0:</li>
          <li>                p_swa.detach().copy_(p_model_)</li>
          <li>            else:</li>
          <li>                p_swa.detach().copy_(self.avg_fn(p_swa.detach(), p_model_,</li>
          <li>                                                 self.n_averaged.to(device)))</li>
          <li>        self.n_averaged += 1</li>
          <li>
          </li>
<li>
          </li>
<li class="selected">@torch.no_grad()</li>
          <li>def update_bn(loader, model, device=None):</li>
          <li>    r"""Updates BatchNorm running_mean, running_var buffers in the model.</li>
          <li>
          </li>
<li>    It performs one pass over data in `loader` to estimate the activation</li>
          <li>    statistics for BatchNorm layers in the model.</li>
          <li>    Args:</li>
          <li>        loader (torch.utils.data.DataLoader): dataset loader to compute the</li>
          <li>            activation statistics on. Each data batch should be either a</li>
          <li>            tensor, or a list/tuple whose first element is a tensor</li>
          <li>            containing data.</li>
      </ol>
    </code></pre>



  </article>

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

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

<p>and I am not exactly sure what is causing this… but I suspect:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/34bdd46e6e6a15dfa1c765c1fb8d9d7bc02dc102/torch/optim/swa_utils.py#L164">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/34bdd46e6e6a15dfa1c765c1fb8d9d7bc02dc102/torch/optim/swa_utils.py#L164" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/34bdd46e6e6a15dfa1c765c1fb8d9d7bc02dc102/torch/optim/swa_utils.py#L164" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/34bdd46e6e6a15dfa1c765c1fb8d9d7bc02dc102/torch/optim/swa_utils.py#L164</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="154" style="counter-reset: li-counter 153 ;">
          <li>    &gt;&gt;&gt; torch.optim.swa_utils.update_bn(loader, model)</li>
          <li>
          </li>
<li>.. note::</li>
          <li>    The `update_bn` utility assumes that each data batch in :attr:`loader`</li>
          <li>    is either a tensor or a list or tuple of tensors; in the latter case it</li>
          <li>    is assumed that :meth:`model.forward()` should be called on the first</li>
          <li>    element of the list or tuple corresponding to the data batch.</li>
          <li>"""</li>
          <li>momenta = {}</li>
          <li>for module in model.modules():</li>
          <li class="selected">    if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):</li>
          <li>        module.running_mean = torch.zeros_like(module.running_mean)</li>
          <li>        module.running_var = torch.ones_like(module.running_var)</li>
          <li>        momenta[module] = module.momentum</li>
          <li>
          </li>
<li>if not momenta:</li>
          <li>    return</li>
          <li>
          </li>
<li>was_training = model.training</li>
          <li>model.train()</li>
          <li>for module in momenta.keys():</li>
      </ol>
    </code></pre>



  </article>

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

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

<p>and that it is not able to find the module in the model? Totally seem clueless <img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=12" title=":sob:" class="emoji" alt=":sob:" loading="lazy" width="20" height="20"></p>
<p>I seem to get better loss with syncbatchnorm and I would like to update the batch stats and perform inference, but the fact that the updates are not working is blocking me <img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=12" title=":sob:" class="emoji" alt=":sob:" loading="lazy" width="20" height="20"></p>
<p>J</p>

smth: <p>do you have a minimal but full script that one can run.<br>
If it’s hanging, and not doing anything, it probably involves a bit more deeper debugging on what’s happening.</p>

James_W: <p>Hi <a class="mention" href="/u/smth">@smth</a></p>
<p>Thank you VERY much for your reply. Well when I read the source, I see this:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/34bdd46e6e6a15dfa1c765c1fb8d9d7bc02dc102/torch/optim/swa_utils.py#L71">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/34bdd46e6e6a15dfa1c765c1fb8d9d7bc02dc102/torch/optim/swa_utils.py#L71" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/34bdd46e6e6a15dfa1c765c1fb8d9d7bc02dc102/torch/optim/swa_utils.py#L71" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/34bdd46e6e6a15dfa1c765c1fb8d9d7bc02dc102/torch/optim/swa_utils.py#L71</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="61" style="counter-reset: li-counter 60 ;">
          <li>Example:</li>
          <li>    &gt;&gt;&gt; # Compute exponential moving averages of the weights and buffers</li>
          <li>    &gt;&gt;&gt; ema_avg = lambda averaged_model_parameter, model_parameter, num_averaged:\</li>
          <li>                        0.1 * averaged_model_parameter + 0.9 * model_parameter</li>
          <li>    &gt;&gt;&gt; swa_model = torch.optim.swa_utils.AveragedModel(model, avg_fn=ema_avg, use_buffers=True)</li>
          <li>
          </li>
<li>.. note::</li>
          <li>    When using SWA with models containing Batch Normalization you may</li>
          <li>    need to update the activation statistics for Batch Normalization.</li>
          <li>    This can be done either by using the :meth:`torch.optim.swa_utils.update_bn`</li>
          <li class="selected">    or by setting :attr:`use_buffers` to `True`. The first approach updates the</li>
          <li>    statistics in a post-training step by passing data through the model. The</li>
          <li>    second does it during the parameter update phase by averaging all buffers.</li>
          <li>    Empirical evidence has shown that updating the statistics in normalization</li>
          <li>    layers increases accuracy, but you may wish to empirically test which</li>
          <li>    approach yields the best results in your problem.</li>
          <li>
          </li>
<li>.. note::</li>
          <li>    :attr:`avg_fn` is not saved in the :meth:`state_dict` of the model.</li>
          <li>
          </li>
<li>.. note::</li>
      </ol>
    </code></pre>



  </article>

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

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

<p>So, I assumed this was the correct way to do it and did it… It works, but now I have a different issue. I want to infer the model on CPU, so I normally do:</p>
<pre><code class="lang-auto">                            traced_script_module = torch.jit.trace(swa_model, (data))
                            traced_script_module.save("./swa_model.tar")
</code></pre>
<p>It complains:</p>
<pre><code class="lang-auto">    raise ValueError("SyncBatchNorm expected input tensor to be on GPU")
ValueError: SyncBatchNorm expected input tensor to be on GPU
</code></pre>
<p>But I want to infer the model on cpu by jit tracing, and I am not sure how I can get around this issue <img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=12" title=":sob:" class="emoji" alt=":sob:" loading="lazy" width="20" height="20"></p>
<p>I presume it is impossible to jit trace with the model on cpu and data on gpu (due to dtype)</p>