DataParallel output differs from its module

username: seongmin
cooked: <p>Hi,</p>
<p>I am trying to do multi-task learning with two classification layers from single shared representation.<br>
I believe, having two nn.Linear layers would not differ from having one nn.Linear layer then split,<br>
but it occurs to me they work differently when used with nn.DataParallel.</p>
<p>Please see the code at the bottom.<br>
&lt;class ‘Split’&gt; has one nn.Linear then splits the output.<br>
&lt;class ‘TwoHeads’&gt; has two nn.Linear layers.</p>
<p>What I do in the main code is to compare the outputs of</p>
<blockquote>
<p>nn.DataParallel(net)</p>
</blockquote>
<p>to</p>
<blockquote>
<p>nn.DataParallel(net).module.</p>
</blockquote>
<p>For Split the output coincides.<br>
For TwoHeads the output differs.</p>
<p>And the code here.</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.nn as nn

class Split(nn.Module):
    def __init__(self):
        super(Split, self).__init__()
        self.linear = nn.Linear(10, 5)

    def forward(self, x):
        out = self.linear(x)
        out1, out2 = torch.split(out, 3, dim=1)
        return out1, out2

class TwoHeads(nn.Module):
    def __init__(self):
        super(TwoHeads, self).__init__()
        self.linear1 = nn.Linear(10, 3)
        self.linear2 = nn.Linear(10, 2)

    def forward(self, x):
        out1 = self.linear1(x)
        out2 = self.linear2(x)
        return out1, out2

if __name__ == '__main__':
    net1 = Split()
    net2 = TwoHeads()

    for net in [net1, net2]:
        net.cuda()
        net = nn.DataParallel(net, list(range(4)))

        with torch.no_grad():
            data = torch.randn(500, 10).cuda()

            out1, out2 = net(data)
            mod1, mod2 = net.module(data)

            print(int(not torch.equal(out1, mod1)), end=' ')
            print(int(not torch.equal(out2, mod2)), end=' ')

            print((out1 - mod1).abs().max(), (out2 - mod2).abs().max())
</code></pre>
<p>For me, the output looks like</p>
<blockquote>
<p>0 0 tensor(0., device=‘cuda:0’) tensor(0., device=‘cuda:0’)<br>
0 1 tensor(0., device=‘cuda:0’) tensor(2.3842e-07, device=‘cuda:0’)</p>
</blockquote>
<p>My questions are:</p>
<ul>
<li>Is it intended?</li>
<li>If it is because the computation graph uses the same variable twice, shall I always avoid ‘branching’ the computation graph?</li>
<li>What would be the correct way to do multi-task learning with DataParallel?</li>
</ul>

username: ptrblck
cooked: <p>The error of <code>1e-7</code> is most likely due to floating point precision (usually you expect the error for FP32 to be in ~1e-6), so it seems your code is working fine.</p>
<pre><code class="lang-python">x = torch.randn(10, 10, 10)
sum1 = x.sum()
sum2 = x.sum(0).sum(0).sum(0)
print(sum1 - sum2)
&gt; tensor(-3.8147e-06)
</code></pre>

username: seongmin
cooked: <p>Thank you so much for your answer and the code.</p>
<p>I firstly thought of it too, and I repeated the same experiment many times and on another machine but<br>
the error still occurs for &lt;class ‘TwoHeads’&gt; at the second output only.</p>
<p>Would it be still most likely the floating point error?</p>

username: ptrblck
cooked: <p>Since the error is that low, I would still assume it’s still due to floating point precision.</p>