My input batch is not splitted on the number of GPUs

username: barakb
cooked: <p>Hi, I have a machine with 4 GPUs, my program went smoothly, I did as always:</p>
<pre><code class="lang-auto">    model = model.cuda()
    model = nn.DataParallel(model) 
</code></pre>
<p>following that I did:</p>
<pre><code class="lang-auto">        images = images.cuda()
        target = target.cuda()
        output = model(images)
</code></pre>
<p>My input batch is 64, 16 on each GPU, everything as planned.</p>
<p>Things are getting wrong when I’m splitting my model to have 2 forwards:</p>
<pre><code class="lang-auto">        midoutput = model.module.forward1(images)
        output = model.module.forward2(midoutput)
</code></pre>
<p>Then although I did the exect same procedure as mentioned above, inside forward1 I’m getting a batch of 64 and not 16 as expected.</p>
<p>when I’m using GPUtil I’m getting:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/b/c/bc4dedae21032f301456500d0864e26506aa0c6c.png" alt="Capture" data-base62-sha1="qROIYg0PlD2I1o8yo1Kdez1Bue8" width="188" height="115"></p>
<p>So, it seems that although I used the same procedure, and I have 4 GPU,  only 1 GPU is working now.</p>
<p>Any Idea?</p>
<p>Thanks a lot!</p>

username: alex.veuthey
cooked: <p>Maybe you splitting the original <code>forward()</code> method in two different ones breaks inheritence and/or autograd in some way?</p>
<p>Why do you need to split it in the first place? I think you could retrieve the <code>midoutput</code> result in the original forward by keeping an intermediate variable and returning it as well as the final output… Something like that:</p>
<pre><code class="lang-python">def forward(x):
    x = self.layer1(x)
    midoutput = x
    x = self.layer2(x)
    return midoutput, x
</code></pre>
<p>and then:</p>
<pre><code class="lang-python">midoutput, output = model(images)
</code></pre>

username: barakb
cooked: <p>I want to loop on one of them , in a way that is not possible using 1 forward function.</p>

username: alex.veuthey
cooked: <p>I see! Maybe you could nest two different modules, keeping only one <code>forward</code> but looping on the submodule’s <code>forward</code> inside it?</p>

username: barakb
cooked: <p>Thanks, but I need to have 2 different forward functions, It’s important for to understand why doing that “breaking” the dataparallel settings, I think it’s also important for the community.</p>

username: alex.veuthey
cooked: <p>Sure, I understand.</p>
<p>Then digging a bit in the doc shows why it acts like this: check the source code for DataParallel <a href="https://pytorch.org/docs/stable/_modules/torch/nn/parallel/data_parallel.html" rel="nofollow noopener">here</a>.</p>
<p>What happens when you call <code>model(images)</code> is that you call the <code>DataParallel</code> module’s <code>forward</code> method, which handles all the parallelism. Inside the <code>forward</code> method, the submodules (your model split in the number of GPUs you specified) are called with <code>return self.module(*inputs[0], **kwargs[0])</code>, which in turn calls the <code>forward</code> method of the module itself… Which is not <code>forward1</code> or <code>forward2</code> in your case, and it doesn’t split it on multiple GPUs.</p>
<p>Hence my recommendations of keeping the original <code>forward</code> method or having two nested modules.</p>
<p>One solution I can imagine is inheriting the <code>DataParallel</code> class, and overriding the <code>forward</code> method with your own looping behaviour, but doing it without breaking anything will probably be a bit complex…</p>

username: barakb
cooked: <p>Thanks for your investigation, I saw the source code and I think you are right.<br>
I will try your solution of two nested modules, can you give a bit more details about this solution?<br>
I can also do something like:</p>
<pre><code class="lang-auto">def forward(self,x,flag):
    if flag:
       # do foraward1 code
    else:
       # do foraward2 code
</code></pre>
<p>And then:</p>
<pre><code class="lang-auto">        images = images.cuda()
        target = target.cuda()
        output = model(images,flag=True)
        output = model(images,flag=False)
</code></pre>
<p>Do you see any drawback to do the above path?</p>

username: alex.veuthey
cooked: <p>I think the second solution should work! I can’t see a reason why it wouldn’t, at least. Did you try it?</p>
<p>For nested modules, here’s a simplified example:</p>
<pre><code class="lang-python">from torch import nn

class SubModule(nn.Module):
    def __init__(self):
        super(SubModule, self).__init__()
        self.layer1 = ...
    def forward(self, x):
        return self.layer1(x)

class Module(nn.Module):
    def __init__(self, submodule):
        super(Module, self).__init__()
        self.submodule = submodule
        self.layer2 = ...
    def forward(self, x):
        x = self.submodule(x)
        x = self.layer2(x)
        return x
</code></pre>