Aggregating the results of Forward / backward hook on nn.DataParallel (multi-GPU)

wwiiiii: <p>Hi, I want to use hook function on DataParallel instance, but I have some uncertain points.<br>
I think below code would work well, but I’m not sure about</p>
<ol>
<li>Do I need to set lock when I access <code>self.target_outputs</code> and <code>self.target_outputs_grad</code> in the hook function?</li>
<li>Is it guaranteed that inputs are scattered to each GPU following original order? For example, if the inputs were [1, 2, 3, 4] and there are 2 GPUs, Do [1, 2] are fed to GPU <span class="hashtag">#1</span> and [3, 4] are fed to GPU <span class="hashtag">#2</span>?</li>
</ol>
<pre><code class="lang-python">class Wrapper:

    def __init__(self, model, target_layer):
        self.model = model
        self.target_layer = target_layer
        self.target_outputs = {}
        self.target_outputs_grad = {}

        def forward_hook(_, __, output):
            self.target_outputs[output.device] = output.detach()

        def backward_hook(_, __, grad_output):
            assert len(grad_output) == 1
            self.target_outputs_grad[grad_output[0].device] = grad_output[0].detach()

        self.target_layer.register_forward_hook(forward_hook)
        self.target_layer.register_backward_hook(backward_hook)
</code></pre>
<p>Any comments would be appreciated! Thanks!</p>

wwiiiii: <p>For Question 1, I found that the pytorch code aggregating output of each (copied) module is using lock.<br>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/parallel/parallel_apply.py#L60" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/parallel/parallel_apply.py#L60" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/master/torch/nn/parallel/parallel_apply.py#L60</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="50" style="counter-reset: li-counter 49 ;">
<li>def _worker(i, module, input, kwargs, device=None):</li>
<li>    torch.set_grad_enabled(grad_enabled)</li>
<li>    if device is None:</li>
<li>        device = get_a_var(input).get_device()</li>
<li>    try:</li>
<li>        with torch.cuda.device(device):</li>
<li>            # this also avoids accidental slicing of `input` if it is a Tensor</li>
<li>            if not isinstance(input, (list, tuple)):</li>
<li>                input = (input,)</li>
<li>            output = module(*input, **kwargs)</li>
<li class="selected">        with lock:</li>
<li>            results[i] = output</li>
<li>    except Exception as e:</li>
<li>        with lock:</li>
<li>            results[i] = e</li>
<li>
</li>
<li>if len(modules) &gt; 1:</li>
<li>    threads = [threading.Thread(target=_worker,</li>
<li>                                args=(i, module, input, kwargs, device))</li>
<li>               for i, (module, input, kwargs, device) in</li>
<li>               enumerate(zip(modules, inputs, kwargs_tup, devices))]</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
<br>
So the answer would be I have to use <code>threading.lock</code>.</p>

wwiiiii: <p>For Question 2, I tracked the <code>scatter</code> of inputs.</p>
<p><a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/parallel/data_parallel.py#L139" rel="nofollow noopener">inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)</a><br>
=&gt; <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/parallel/data_parallel.py#L150" rel="nofollow noopener">scatter_kwargs(inputs, kwargs, device_ids, dim=self.dim)</a><br>
=&gt; <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/parallel/scatter_gather.py#L13" rel="nofollow noopener">Scatter.apply(target_gpus, None, dim, obj)</a><br>
=&gt; <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/parallel/_functions.py#L89" rel="nofollow noopener">comm.scatter(input, target_gpus, chunk_sizes, ctx.dim, streams)</a><br>
=&gt; <a href="https://github.com/pytorch/pytorch/blob/master/torch/cuda/comm.py#L142" rel="nofollow noopener">tuple(torch. <em>C.</em> scatter(tensor, devices, chunk_sizes, dim, streams))</a><br>
=&gt; <a href="https://github.com/pytorch/pytorch/blob/277b637811abd95f248755bff61c501aa210a040/torch/csrc/cuda/python_comm.cpp#L49" rel="nofollow noopener">scatter(tensor, devices, chunk_sizes, dim, streams);</a><br>
=&gt; <a href="https://github.com/pytorch/pytorch/blob/277b637811abd95f248755bff61c501aa210a040/torch/csrc/cuda/comm.cpp#L148" rel="nofollow noopener">tensor.chunk(/ <em>chunks=</em> /devices.size(), / <em>dim=</em> /dim);</a><br>
=&gt; <a href="https://github.com/pytorch/pytorch/blob/57ec8f111f912c2b3c1d40ab4ac3ba36a93b3182/aten/src/ATen/native/TensorShape.cpp#L75" rel="nofollow noopener">self.split(split_size, dim);</a><br>
=&gt; <a href="https://github.com/pytorch/pytorch/blob/57ec8f111f912c2b3c1d40ab4ac3ba36a93b3182/aten/src/ATen/native/TensorShape.cpp#L375" rel="nofollow noopener">self.narrow(dim, i * split_size, length);</a></p>
<p>It seems inputs are scattered to GPU in its original order, so aggregating the outputs of each GPU in GPU # order would be matched with original inputs. Indeed, I can find it at  <code>gather</code> of <code>nn.DataParallel</code>.</p>

PangWong: <p>So would you please show us what your code looks like finally.<br>
I also tried to get intermedia outputs by forward hook using multi-gpus but a weird thing happened.</p>
<ul>
<li>In <strong>init</strong>, I initialize <code>self.target_output</code> as None.</li>
<li>In hook function, I print <code>self.target_output</code> which is <code>not None</code> when hook function is excuted.</li>
<li>But after <code>self.model.forward()</code> excuted, <code>self.target_output</code>  turn back to be None.</li>
</ul>
<p>One of two forwards output None  which is vary weird.</p>
<p>I implement this using <code>data_parellel</code> function and I regard <code>class Wrapper</code> as a NN Module, which return <code>self.target_output</code>. So <code>data_parallell</code> will scatter and gather all the output of <code>Wrapper</code>.</p>
<p>code looks like:</p>
<pre><code class="lang-python">import torch
from torchvision.models.vgg import vgg19

class Wrapper(torch.nn.Module):
    def __init__(self, model):
        super(Wrapper, self).__init__()
        self.model = model
        self.target_outputs = None

        def forward_hook(_, __, output):
            self.target_outputs = output.detach()
        self.model.features[2].register_forward_hook(forward_hook)

    def forward(self,input):
        self.model(input)
        return self.target_outputs


model = vgg19()
model = model.cuda(4)
wrapper = Wrapper(model)

devices = [4, 5]

inputs = torch.randn(60,3,224,224).fill_(0).cuda(4)
out = torch.nn.parallel.data_parallel(wrapper, inputs, devices)
print("first forward:   ", out)
inputs = torch.randn(60,3,224,224).fill_(1).cuda(4)
out = torch.nn.parallel.data_parallel(wrapper, inputs, devices)
print("second forward: ", out.shape)
</code></pre>
<p>output is:</p>
<pre><code class="lang-auto">first forward:  None
second forward: (60, 64, 224, 224)

</code></pre>
<p>By compare the result with single-gpu result, I found the result of second forward is indeed the result of first forward .</p>

PangWong: <p>There is a solution to get intermedia output by forward hooks using multi-gpus in <a href="https://discuss.pytorch.org/t/weird-output-of-forward-hook-when-use-multi-gpus/29437/2?u=pangwong">this post</a> although not so elegant. It dose work in my test.<br>
But the weird phenomenon explained in above post  still confuses me.</p>

qiuyang: <p><a class="mention" href="/u/wwiiiii">@wwiiiii</a> Hi,do you know how  does Scatter.apply fuction call Scatter.forward function?I am confused about this.</p>

Adria_Ciurana: <p>This answer may come late, but I hope it will serve someone. Or if someone finds a better way.</p>
<p>I solved using thread identificator. It is not the best way, but it has worked for me. This way you will avoid blocking two processes.</p>
<pre><code class="lang-auto">import torch
from torchvision.models.vgg import vgg19
import threading
from collections import defaultdict

class Wrapper(torch.nn.Module):
    def __init__(self, model):
        super(Wrapper, self).__init__()
        self.model = model
        self.target_outputs = defaultdict(lambda: None)

        def forward_hook(_, __, output):
            self.target_outputs[threading.get_native_id()] = output.detach()
        self.model.features[2].register_forward_hook(forward_hook)

    def forward(self,input):
        self.model(input)

        thread_id = threading.get_native_id()
        result = self.target_outputs[thread_id]
        del self.target_outputs[thread_id]

        return result


model = vgg19()
model = model.cuda(4)
wrapper = Wrapper(model)

devices = [4, 5]

inputs = torch.randn(60,3,224,224).fill_(0).cuda(4)
out = torch.nn.parallel.data_parallel(wrapper, inputs, devices)
print("first forward:   ", out)
inputs = torch.randn(60,3,224,224).fill_(1).cuda(4)
out = torch.nn.parallel.data_parallel(wrapper, inputs, devices)
print("second forward: ", out.shape)
</code></pre>

shbmgpta: <p>Hi,<br>
Thanks for the amazing solution. It works correctly for the forward pass.<br>
I am using forward along with backward and doing certain operations. So, each time, when the backward is called, it creates a new thread and the output of backward is saved on that thread. My problem is how do I get the backward thread ID for the corresponding forward thread ID.</p>
<pre><code>self.forward_hook = self.target_layer.register_forward_hook(self.hook_fn_act)
self.backward_hook = self.target_layer.register_backward_hook(self.hook_fn_grad)

def hook_fn_act(self, module, input, output):
    self.activations[threading.get_native_id()] = output.detach()

def hook_fn_grad(self, module, grad_input, grad_output):
    self.gradients[threading.get_native_id()] = grad_output[0].detach()

def foo():
    &lt;Certain operations using the self.activations and self.gradients&gt;
    &lt;Problem is that thread id for the activations and gradients is not same, so how to know which 
      thread ID of activations maps to which one of gradients&gt;</code></pre>

frostinassiky: <p>Thanks for the great solution.</p>
<p>For python3.7, we need to use <code>threading.get_ident()</code> because <code>get_native_id()</code> is only defined in python3.8.</p>
<p>The following code works for me.</p>
<pre><code class="lang-auto">In [4]: import torch
   ...: import threading
   ...: from collections import defaultdict
   ...: from torchvision.models.vgg import vgg19
   ...: 
   ...: class Wrapper(torch.nn.Module):
   ...:     def __init__(self, model):
   ...:         super(Wrapper, self).__init__()
   ...:         self.model = model
   ...:         self.target_outputs = {}
   ...: 
   ...:         def forward_hook(_, __, output):
   ...:             self.target_outputs[threading.get_ident()] = output.detach()
   ...:         self.model.features[2].register_forward_hook(forward_hook)
   ...: 
   ...:     def forward(self,input):
   ...:         self.model(input)
   ...:         thread_id = threading.get_ident()
   ...:         result = self.target_outputs[thread_id]
   ...:         del self.target_outputs[thread_id]
   ...:         return result
   ...: 
   ...: 
   ...: model = vgg19()
   ...: model = model.cuda(0)
   ...: wrapper = Wrapper(model)
   ...: 
   ...: devices = [0, 1]
   ...: 
   ...: inputs = torch.randn(60,3,224,224).fill_(0).cuda(1)
   ...: out = torch.nn.parallel.data_parallel(wrapper, inputs, devices)
   ...: print("first forward:   ", out.shape)
   ...: inputs = torch.randn(60,3,224,224).fill_(1).cuda(1)
   ...: out = torch.nn.parallel.data_parallel(wrapper, inputs, devices)
   ...: print("second forward: ", out.shape)
</code></pre>
<p>The result is</p>
<pre><code class="lang-auto">first forward:    torch.Size([60, 64, 224, 224])
second forward:  torch.Size([60, 64, 224, 224])
</code></pre>