DataParallel does not concat outputs from multi-gpu

username: pyxiea
cooked: <p>A similar problem on stackoverflow <a href="https://stackoverflow.com/questions/57607680/how-does-torch-nn-dataparallel-change-the-output-size" rel="noopener nofollow ugc">is here</a>, but no answer is useful.</p>
<p>As <a href="https://pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html" rel="noopener nofollow ugc">this tutorial</a> shown, the output of multi-gpus will be concatenated on the dimension 0, but I don’t know why does it not work in my code.</p>
<pre><code class="lang-python">
    model = T2T(......)  # T2T is a sub class of nn.Module
    if torch.cuda.device_count() &gt; 1:
        print("Using", torch.cuda.device_count(), "GPUs!")
        model = torch.nn.DataParallel(model)
    model = model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
    train_loader = DataLoader(......)
    ......
    epoch_steps = len(train_loader)
    STEPS = train_epochs * epoch_steps
    train_loader = cycle(train_loader)
    ......
    for step in range(STEPS):
        batch_data = next(train_loader)
        optimizer.zero_grad()
        labels = batch_data['labels'].to(device)
        ......
        probs = model(......)
        # calculate loss using labels and probs
</code></pre>
<p>When calculating loss, I got an error, the batch size of <code>labels</code> and <code>probs</code> are not the same. The shape of <code>labels</code> is [64,…] and the shape of <code>probs</code> is [32,…]. I am using 2 gpus, so I guess the output of multi-gpus are not concatenated.</p>
<p>Any idea how to fix this? Thanks in advance!</p>

username: ptrblck
cooked: <p>What is the shape of the input? The <code>nn.DataParallel</code> model would split the input in <code>dim0</code> and concatenate it afterwards. If the input has a batch size of 32, the output would have the same shape.</p>

username: pyxiea
cooked: <pre><code class="lang-python">probs = model(inputrulelist, syn_inputrulelist, tree_path_vec, rule_mask, syn_rule_mask, inputrulelistnode, syn_inputrulelistnode, inputrulelistson, syn_inputrulelistson, sequence_mask, treemask, syn_treemask, path_lens)
</code></pre>
<p>Sorry for the late reply, I have already fixed the problem but the actual reason remains unknown. The model has many params as above, but the first dimensions of them are all <code>batch_size</code>, and the current value is <code>64</code>.</p>
<p>One exception is the <code>sequence_mask</code> tensor, it has shape <code>[1, n, n]</code> and the elements on and below the main diagonal are all zeros (it is used to do the sequence masking in multi-head attention). The <code>batch_size</code> of it is set to <code>1</code> because I have noticed that <code>DataParallel</code> model will split the input in <code>dim0</code>, and I hope all devices will get the same <code>sequence_mask</code>.</p>
<p><strong>After I change the first dimension of <code>sequence_mask</code> to <code>batch_size</code>(64 for now), the problem is fixed.</strong> Specifically, I use <code>tensor.repeat(batch_size, 1, 1)</code> instead of <code>tensor.unsqueeze(0)</code> to generate the <code>batch_size</code> dimension as before. After that, the first dimension of the output of the <code>DataParallel</code> model (<code>probs</code> in the code) becomes <code>64</code> instead of <code>32</code>. I don’t know the exact reason due to my unfamiliarity of <code>DataParallel</code>, it would be so helpful if you could tell me, thanks a lot.</p>

username: ptrblck
cooked: <p>Your description is right. <code>nn.DataParallel</code> expects all input tensors in the shape <code>[batch_size, *]</code> and will split them in <code>dim0</code>. If you’ve applied broadcasting in the past, repeating the tensor sounds reasonable.</p>

username: pyxiea
cooked: <p>But I still don’t know why can I only get output from one gpu (instead of the concatenation of outputs from all gpus) if I pass a tensor with shape [1, *]</p>

username: ptrblck
cooked: <p>This tensor won’t be split, as <code>dim0</code> has only a single sample and thus only a single GPU will execute this tensor.<br>
If you want to make sure each GPU gets a chunk of the input tensor, make sure <code>dim0</code> has at least a size of <code>num_gpus</code>.</p>

username: pyxiea
cooked: <p>Oh I get it, thank you for your time!</p>