DataParallel breaks the forward() method

username: badmax
cooked: <p>Hi,</p>
<p>When I wrap my model in DataParallel my forward() method no longer gets input of the right size. For example</p>
<pre><code class="lang-auto">class Net(nn.Module):
    def __init__(self):
        super().__init__()
        
    def forward(self, x, y):
        return x, y

net = Net()

device = torch.device('cuda')
net = nn.DataParallel(net)

net = net.to(device)

x = torch.ones((1024, 3, 100, 100))
y = torch.ones((1, 3, 100, 100))

x_out, y_out = net(x, y)

print(x_out.shape)
#torch.Size([128, 3, 100, 100])

</code></pre>
<p>I have 8 GPUs so it seems like I only get back the result from one of them.  If I change the first dimension of <code>y</code> to match <code>x</code> I get back the original result, as expected.</p>

username: eqy
cooked: <p>My guess is something is strange about trying to coalesce <code>y</code> across the “batch” dimension when it is getting passed back. Does working around this by removing the first dimension of <code>y</code> (e.g., <code>y = y.squeeze(0)</code>) before passing it to the net work?</p>

username: badmax
cooked: <p>No, in that case I get <code>torch.Size([384, 3, 100, 100])</code></p>

username: wayi
cooked: <p><a class="mention" href="/u/vitalyfedyunin">@VitalyFedyunin</a> for DataParallel question.</p>

username: ptrblck
cooked: <p>If you add some debug statements to the <code>forward</code> method you’ll see that due to the second input, the processing fails, since it cannot be chunked into 8 parts:</p>
<pre><code class="lang-python">class Net(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x, y):
        print('x: {}, {}'.format(x.device, x.shape))
        print('y: {}, {}'.format(y.device, y.shape))
        return x, y
</code></pre>
<p>Your setup:</p>
<pre><code class="lang-python">x = torch.ones((1024, 3, 100, 100))
y = torch.ones((1, 3, 100, 100))

x_out, y_out = net(x, y)

print(x_out.shape)

x: cuda:0, torch.Size([128, 3, 100, 100])
y: cuda:0, torch.Size([1, 3, 100, 100])
torch.Size([128, 3, 100, 100])
</code></pre>
<p>I’m currently unsure what the expected behavior is, as it currently seems to fall back to the smallest possible splitting (I would assume an error would be raised).</p>
<p>If you are using tensors with the same shape in <code>dim0</code>, it’s working as expected:</p>
<pre><code class="lang-python">x = torch.ones((1024, 3, 100, 100))
y = torch.ones((1024, 3, 100, 100))

x_out, y_out = net(x, y)

print(x_out.shape)

x: cuda:0, torch.Size([128, 3, 100, 100])
y: cuda:0, torch.Size([128, 3, 100, 100])
x: cuda:1, torch.Size([128, 3, 100, 100])
y: cuda:1, torch.Size([128, 3, 100, 100])
x: cuda:2, torch.Size([128, 3, 100, 100])
x: cuda:3, torch.Size([128, 3, 100, 100])
y: cuda:2, torch.Size([128, 3, 100, 100])
x: cuda:4, torch.Size([128, 3, 100, 100])
x: cuda:5, torch.Size([128, 3, 100, 100])
y: cuda:4, torch.Size([128, 3, 100, 100])
y: cuda:5, torch.Size([128, 3, 100, 100])
x: cuda:6, torch.Size([128, 3, 100, 100])
x: cuda:7, torch.Size([128, 3, 100, 100])
y: cuda:6, torch.Size([128, 3, 100, 100])
y: cuda:7, torch.Size([128, 3, 100, 100])
y: cuda:3, torch.Size([128, 3, 100, 100])
torch.Size([1024, 3, 100, 100])
</code></pre>