Multi GPU backwards hook on wrong device

username: pytorcher
cooked: <p>I have a backward hook function with ‘newLayer.register_backward_hook(hook_function)’ where I do not know how to control the inputs to it.  The function contains a line like.</p>
<pre><code class="lang-auto">def hook_function(self, grad_input, grad_output):
    self.average = self.average * 0.99 + grad_output[0].sum((0,2,3)) * 0.01
</code></pre>
<p>This results in:</p>
<pre><code class="lang-auto">RuntimeError: expected device cuda:0 but got device cuda:1
(Pdb) self.average
tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], device='cuda:0',
       dtype=torch.float64)
(Pdb) grad_output[0].sum((0,2,3))
tensor([ 0.0508,  0.0492,  0.0512,  0.0487,  0.0517,  0.0483,  0.0522,  0.0479,
        -0.1974, -0.2026], device='cuda:1', dtype=torch.float64)

</code></pre>
<p>I need to know either how to define my self.average better or if i am doing something wrong somewhere else and it is grad_output that is wrong.</p>
<p>the above is with</p>
<pre><code class="lang-auto">self.average = torch.Tensor(out_channels).zero_().to(gf.device).double()
</code></pre>
<p>I have also tried</p>
<pre><code class="lang-auto">self.average = nn.Parameter(torch.Tensor(out_channels).zero_().to(gf.device).double())
</code></pre>
<p>which results in</p>
<pre><code class="lang-auto">TypeError: cannot assign 'torch.cuda.DoubleTensor' as parameter 'normalPassAverageD' (torch.nn.Parameter or None expected)
</code></pre>
<p>at the same location</p>

username: ptrblck
cooked: <p>I don’t know how <code>self.average</code> is initialized, but would assume this should work:</p>
<pre><code class="lang-python">self.average = self.average.to(grad_output[0].device) * 0.99 + grad_output[0].sum((0,2,3)) * 0.01
</code></pre>
<p>Could you check it and see, if you are still getting an error?<br>
In that case, could you post a small code snippet to reproduce this issue?</p>

username: pytorcher
cooked: <p>That worked!  Thanks!</p>
<p>But this is all within a DataParallel.  Isn’t that supposed to make it play nice with multiple devices?  I would have thought calling .to(a specific device) would limit it to only using that device.  Maybe I don’t understand how DataParallel works entirely.</p>

username: ptrblck
cooked: <p>It’s hard to tell what’s creating the issue without seeing the code, but generally you are right. <code>nn.DataParallel</code> should take care of pushing the parameters and buffers to the right devices. However, if  <code>self.average</code> was defined as at tensor, you would run into this error.</p>

username: pytorcher
cooked: <p>It was in fact defined as a tensor.  Is there another way I should define self tensors to be sure they are handled by DataParallel?</p>

username: ptrblck
cooked: <p>If you don’t need to update this tensor (i.e. no gradients should be calculated for it), you should define it as a buffer via: <code>self.register_buffer('average', torch.tensor(...))</code>. This would make sure that <code>nn.DataParallel</code> will push this buffer to the corresponding device. In your <code>forward</code> method you can access it via <code>self.average</code>.</p>

username: pytorcher
cooked: <p>Cool!  That worked for all of my tensors that are having their values be modified without the .to().  However, later in the same function there is one more tensor I am just saving directly with</p>
<pre><code class="lang-auto">self.current = grad_output[0].detach()
</code></pre>
<p>This works fine on one GPU but when running on multiple it tells me 'AttributeError: (module) object has no attribute ‘current’ '.  I imagine I could do a .copy_() or something, but I won’t actually know the shape until the network has been run on an image.  Thoughts?  I can also start a new thread if that would be appropriate.</p>

username: ptrblck
cooked: <p>Manually manipulating model attributes when running a data parallel approach is a bit tricky, since each device uses a replica of the model and your changes could be lost.<br>
Feel free to create a new topic and describe your use case further.</p>

username: pytorcher
cooked: <aside class="quote" data-post="1" data-topic="93967">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/p/7ea924/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/initializing-a-member-tensor-after-creation-with-dataparallel-repost/93967">Initializing a member tensor after creation with DataParallel (repost)</a> 
  </div>
  <blockquote>
    I have a member tensor that is created/saved during a backward hook. Hook looks like: 
def saveGrad(self, grad_input, grad_output):
        self.currentGrad = grad_output[0].detach()

This works fine on one GPU but when running on multiple it tells me 'AttributeError: (module) object has no attribute ‘currentGrad’ '. I imagine I could do a .copy_() or something, but I won’t actually know the shape until the network has been run on an image.
  </blockquote>
</aside>


username: pytorcher
cooked: <p>Finally back to this after debugging other things.  Apparently the .to() method made it run but it doesn’t get the same values which I imagine is because its not actually keeping the numbers the same.  I’m using a registered buffer as well.  Started a thread with all my current code, please take a look when you get a chance <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"> <a href="https://discuss.pytorch.org/t/multi-gpu-hook-not-correctly-filling-buffer/96046" class="inline-onebox">Multi GPU Hook not correctly filling buffer</a></p>

username: pytorcher
cooked: <p>Have another thread now that includes a full code sample.</p><aside class="quote quote-modified" data-post="5" data-topic="101305">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/p/7ea924/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/multi-gpu-hook-not-correctly-doing-filling-buffer/101305/5">Multi GPU Hook not correctly doing filling buffer</a> 
  </div>
  <blockquote>
    Sorry, I kept a bit of extra code since I started with <a href="https://github.com/pytorch/examples/blob/master/mnist/main.py" rel="noopener nofollow ugc">https://github.com/pytorch/examples/blob/master/mnist/main.py</a>. and thought it would be easy to compare to that.  I have further reduced the code below to only have the exact things to reproduce this problem.  My network is 1 layer, the ‘main’ just inititalizes things and runs backprop on one random input.  The only thing that is unique is my saveAverageD function and 2 custom modules.  The custom modules and saveAverageD function were made b…
  </blockquote>
</aside>