Parameter created with register_buffer is not updated with MultipleGPUs

Cosimo_Rulli: <p>I am creating my own module extending Conv2d. Inside the constructor I define my own “my_param” module, which I use to store some information on the weights during forward. Then I pass this module tothe  F.Conv2d function, in place of self.weight as in classic convolution.</p>
<pre><code class="lang-auto">class MyConv(nn.Conv2d):
    def __init__(...):
...
         self.register_buffer("my_param", torch.zeros(self.weight.shape))
    

    def forward(self, input):
        self.my_param = special_op(self.weight) 

        output = F.conv2d(input, self.my_param, self.bias,
                          self.stride, self.padding,
                          self.dilation, self.groups)
        return output
    
</code></pre>
<p>When using this class with 1 GPU (even with DataParallel), self.my_param correctly stores the ouput of special_op. Instead, with multiple GPUs, self.my_param remains a tensor of zeros.</p>
<p>Is there any workaround? What  am I missing?</p>

tom: <p>As you observe, by default, the buffers are local to the GPU.</p>
<p>If you wanted to synchronize your buffer across GPUs, you need to provide the structure to do that, which is a bit tricky in general.</p>
<p>SyncBatchNorm might be some inspiration (note that it also imports a SyncBatchNorm Function).</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/e86d8323cb872a835ba0845e3519079f25e204fc/torch/nn/modules/batchnorm.py#L728">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/e86d8323cb872a835ba0845e3519079f25e204fc/torch/nn/modules/batchnorm.py#L728" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/e86d8323cb872a835ba0845e3519079f25e204fc/torch/nn/modules/batchnorm.py#L728" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/e86d8323cb872a835ba0845e3519079f25e204fc/torch/nn/modules/batchnorm.py#L728</a></h4>



  <pre class="onebox">    <code class="lang-py">
      <ol class="start lines" start="718" style="counter-reset: li-counter 717 ;">
          <li># If buffers are not to be tracked, ensure that they won't be updated</li>
          <li>running_mean = (</li>
          <li>    self.running_mean if not self.training or self.track_running_stats else None</li>
          <li>)</li>
          <li>running_var = (</li>
          <li>    self.running_var if not self.training or self.track_running_stats else None</li>
          <li>)</li>
          <li>
          </li>
<li># Don't sync batchnorm stats in inference mode (model.eval()).</li>
          <li>need_sync = (bn_training and self.training)</li>
          <li class="selected">if need_sync:</li>
          <li>    process_group = torch.distributed.group.WORLD</li>
          <li>    if self.process_group:</li>
          <li>        process_group = self.process_group</li>
          <li>    world_size = torch.distributed.get_world_size(process_group)</li>
          <li>    need_sync = world_size &gt; 1</li>
          <li>
          </li>
<li># fallback to framework BN when synchronization is not necessary</li>
          <li>if not need_sync:</li>
          <li>    return F.batch_norm(</li>
          <li>        input,</li>
      </ol>
    </code>
  </pre>


  </article>

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

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

<p>Best regards</p>
<p>Thomas</p>