DataParallel doesn't work when calling model.module.some_attribute

gollum: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/6/6cdde65ed7679ebfb8630f288f6dd57209fddc85.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/6cdde65ed7679ebfb8630f288f6dd57209fddc85" title="47%20PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/6/6cdde65ed7679ebfb8630f288f6dd57209fddc85_2_680x500.jpeg" alt="47%20PM" data-base62-sha1="fx4Z4NocxYy6DwVCjRMQUKXobOd" width="680" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/6/6cdde65ed7679ebfb8630f288f6dd57209fddc85_2_680x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/2X/6/6cdde65ed7679ebfb8630f288f6dd57209fddc85_2_1020x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/2X/6/6cdde65ed7679ebfb8630f288f6dd57209fddc85_2_1360x1000.jpeg 2x" data-dominant-color="F3F3F3"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">47%20PM</span><span class="informations">1776×1304 281 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/0/0fa2b6a278f7ceeb6b2889d18c5ef1aa0d07e2e2.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/0fa2b6a278f7ceeb6b2889d18c5ef1aa0d07e2e2" title="26%20PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/0/0fa2b6a278f7ceeb6b2889d18c5ef1aa0d07e2e2_2_677x500.jpeg" alt="26%20PM" data-base62-sha1="2ejLZ2xYedrPdbZQ3Y1c0lYHVtM" width="677" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/0/0fa2b6a278f7ceeb6b2889d18c5ef1aa0d07e2e2_2_677x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/2X/0/0fa2b6a278f7ceeb6b2889d18c5ef1aa0d07e2e2_2_1015x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/2X/0/0fa2b6a278f7ceeb6b2889d18c5ef1aa0d07e2e2_2_1354x1000.jpeg 2x" data-dominant-color="F2F2F2"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">26%20PM</span><span class="informations">1800×1328 327 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
So in the two images there are two different models <code>model</code> and <code>model_p</code> both being wrapped under <code>nn.DataParallel</code>. But in <code>model</code> when calling some attribute <code>fit</code> using the <code>model.module</code> method, I’m unable utilize the two GPUs I originally wanted to parallelize my model upon.  i.e <code>model</code> doesn’t split the dim=0 batch_first dimension into two equal halves for putting it onto two devices as can be seen from the print statements.<br>
Ps. I am very new to using DataParallel and wanted to use something like this. i.e What I actually want is, to call <code>model.module.fit</code> in my training loop with the args as the inputs from my dataloader and in this fit attribute ultimately will makes a call to the forward method of the class model.</p>
<p>But this whole thing doesn’t seem to parallelize and utilize the two GPUs which the model_p could without any fit function and a direct call to forward internally.</p>
<p>I’ve added the <a href="https://drive.google.com/open?id=1yyipI2Mnud8X2Gc8jgoEj3V1v-AOJPIi" rel="noopener nofollow ugc">link</a> to the notebook which was run with CUDA_VISIBLE_DEVICES=0,1</p>
<p>What should I change?<br>
Thanks!</p>
<pre><code class="lang-auto">class Model(nn.Module):
    # Our model


    def __init__(self, input_size, output_size):
        super(Model, self).__init__()
        self.fc = nn.Linear(input_size, output_size)

    def forward(self, input):
        output = self.fc(input)
        return output
    
    def fit(self, input):
        output = self.forward(input)
        print("\tIn Model: input size", input.size(),
              "output size", output.size())
        return output

model = Model(input_size, output_size)
if torch.cuda.device_count() &gt; 1:
  print("Let's use", torch.cuda.device_count(), "GPUs!")
  # dim = 0 [30, xxx] -&gt; [10, ...], [10, ...], [10, ...] on 3 GPUs
  model = nn.DataParallel(model)

model.to(device)

for data in rand_loader:
    input = data.to(device)
    output = model.module.fit(input)
    print("Outside: input size", input.size(),"output_size", output.size())

#############################CASE 2############################
class ModelParallel(nn.Module):
    # Our model

    def __init__(self, input_size, output_size):
        super(ModelParallel, self).__init__()
        self.fc = nn.Linear(input_size, output_size)

    def forward(self, input):
        output = self.fc(input)
        print("\tIn Model: input size", input.size(),
              "output size", output.size())

        return output
    
model_p = ModelParallel(input_size, output_size)
if torch.cuda.device_count() &gt; 1:
  print("Let's use", torch.cuda.device_count(), "GPUs!")
  # dim = 0 [30, xxx] -&gt; [10, ...], [10, ...], [10, ...] on 3 GPUs
  model_p = nn.DataParallel(model_p)
#   model.module.fit = nn.DataParallel(model.module.fit)

model_p.to(device)

for data in rand_loader:
    input = data.to(device)
    output = model_p(input)
    print("Outside: input size", input.size(),"output_size", output.size())
</code></pre>

jerinphilip: <p>DataParallel splits GPUs using its custom forward function and is implemented as a  wrapper rather than a subclass which overrides the model’s <code>forward</code>. When you’re calling fit, you’re calling the <code>forward()</code> associated with the model and not the one wrapped around <code>DataParallel</code>. Hence it will only use a single gpu, as the scatter gather in <code>DataParallel.forward(...)</code> is never called.</p>
<p>From <a href="https://pytorch.org/docs/stable/_modules/torch/nn/parallel/data_parallel.html" rel="nofollow noopener">docs</a>:</p>
<pre><code class="lang-python">    def forward(self, *inputs, **kwargs):
        if not self.device_ids:
            return self.module(*inputs, **kwargs)

        for t in chain(self.module.parameters(), self.module.buffers()):
            if t.device != self.src_device_obj:
                raise RuntimeError("module must have its parameters and buffers "
                                   "on device {} (device_ids[0]) but found one of "
                                   "them on device: {}".format(self.src_device_obj, t.device))

        inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
        if len(self.device_ids) == 1:
            return self.module(*inputs[0], **kwargs[0])
        replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
        outputs = self.parallel_apply(replicas, inputs, kwargs)
        return self.gather(outputs, self.output_device)
</code></pre>

gollum: <p>Thanks for the answer <a class="mention" href="/u/jerinphilip">@jerinphilip</a> ! Makes sense why it only uses one gpu in the case of <code>fit</code><br>
So is there anything I can do which can help me in parallelizing the fit method? Or the only way to parallelize a model is to call the forward from dataparallel wrapped model itself?</p>

jerinphilip: <p>You can use a flag keyword argument inside forward, noticing that the two functions don’t differ by much. I tried to switch member functions using a flag and the following segment worked for me:</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/jerinphilip/MaskGAN.pytorch/blob/cdb2a7aa87826464f79273976286e90fbd3845dc/mgan/modules/distributed_model.py#L58-L70" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/jerinphilip/MaskGAN.pytorch/blob/cdb2a7aa87826464f79273976286e90fbd3845dc/mgan/modules/distributed_model.py#L58-L70" target="_blank" rel="nofollow noopener">jerinphilip/MaskGAN.pytorch/blob/cdb2a7aa87826464f79273976286e90fbd3845dc/mgan/modules/distributed_model.py#L58-L70</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="58" style="counter-reset: li-counter 57 ;">
<li>def forward(self, *args, **kwargs):</li>
<li>    if 'ppl' not in kwargs:</li>
<li>        kwargs['ppl'] = False</li>
<li>
</li>
<li>    if kwargs['tag'] == 'g-step':</li>
<li>        if self.pretrain:</li>
<li>            return self._gstep_pretrain(*args, ppl_compute=kwargs['ppl'])</li>
<li>        else:</li>
<li>            return self._gstep(*args, ppl_compute=kwargs['ppl'])</li>
<li>    elif kwargs['tag'] == 'c-step':</li>
<li>        return self._cstep(*args)</li>
<li>
</li>
<li>    return self._dstep(*args, real=kwargs['real'])</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>In my case I’m switching whether to use generator, discriminator or critic in an GAN-Actor Critic setup. I’m using <code>tag</code> here to control which sub-model’s forward is being called.</p>
<p>You can see <code>scatter</code>'s source code below to understand how <code>args</code> and <code>kwargs</code> are replicated along workers, in case there’s any confusion, which at the time I had:</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/4e3c97a0be5c1bba04928de6abbdad31169e62ee/torch/nn/parallel/scatter_gather.py#L5-L31" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/4e3c97a0be5c1bba04928de6abbdad31169e62ee/torch/nn/parallel/scatter_gather.py#L5-L31" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/4e3c97a0be5c1bba04928de6abbdad31169e62ee/torch/nn/parallel/scatter_gather.py#L5-L31</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="5" style="counter-reset: li-counter 4 ;">
<li>def scatter(inputs, target_gpus, dim=0):</li>
<li>    r"""</li>
<li>    Slices tensors into approximately equal chunks and</li>
<li>    distributes them across given GPUs. Duplicates</li>
<li>    references to objects that are not tensors.</li>
<li>    """</li>
<li>    def scatter_map(obj):</li>
<li>        if isinstance(obj, torch.Tensor):</li>
<li>            return Scatter.apply(target_gpus, None, dim, obj)</li>
<li>        if isinstance(obj, tuple) and len(obj) &gt; 0:</li>
<li>            return list(zip(*map(scatter_map, obj)))</li>
<li>        if isinstance(obj, list) and len(obj) &gt; 0:</li>
<li>            return list(map(list, zip(*map(scatter_map, obj))))</li>
<li>        if isinstance(obj, dict) and len(obj) &gt; 0:</li>
<li>            return list(map(type(obj), zip(*map(scatter_map, obj.items()))))</li>
<li>        return [obj for targets in target_gpus]</li>
<li>
</li>
<li>    # After scatter_map is called, a scatter_map cell will exist. This cell</li>
<li>    # has a reference to the actual function scatter_map, which has references</li>
<li>    # to a closure that has a reference to the scatter_map cell (because the</li>
</ol></code></pre>

  This file has been truncated. <a href="https://github.com/pytorch/pytorch/blob/4e3c97a0be5c1bba04928de6abbdad31169e62ee/torch/nn/parallel/scatter_gather.py#L5-L31" target="_blank" rel="nofollow noopener">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


gollum: <p>Thanks a tonne for helping me out! <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=9" title=":smile:" class="emoji" alt=":smile:"><br>
I found a simple way to change my code and use the parallel functionality in my forward.<br>
Everything working as expected now.</p>

Sam_Sam: <p>Hi, I met the same problem and want to use multi gpus even with model.module.predict, predict is a part defined in the model class. So, could you tell me the simple way you <a class="mention" href="/u/gollum">@gollum</a>  found? Thanks!</p>