Is the loss function paralleled when using DataParallel?

el3ment: <p>When using <a href="http://pytorch.org/docs/_modules/torch/nn/parallel/data_parallel.html" rel="nofollow noopener">DataParallel</a> to wrap my module, do I need to do anything to also parallelize the loss functions?</p>
<p>For example, let’s say that I have large batch size and large output tensors to compute MSE against a target. This operation would benefit from splitting the batch across multiple GPUs, but I’m not sure if the following code does that:</p>
<pre><code>model = MyModule()
model = nn.parallel.DataParallel(model, device_ids=range(args.number_gpus))
model.cuda()

output = model(data)
criterion = nn.MSELoss()
criterion.cuda()
loss = criterion(output, target)
loss.backward()
optimizer.step()
</code></pre>
<p>This is a simplification based on <a href="https://github.com/pytorch/examples/blob/0984955bb8525452d1c0e4d14499756eae76755b/imagenet/main.py" rel="nofollow noopener">imagenet example</a>.</p>
<p>In my case, I have a much bigger custom loss module that includes some calls to a VGG network to estimate perceptual loss, and I’m not sure if I am maximizing performance. I tried computing loss as part of the forward function in <code>MyModule</code>, but this led to recursion errors during the backward step.</p>

smth: <p>you can wrap the loss function inside a DataParallel too if you’d like.</p>

el3ment: <p>Would that result in one unnecessary scatter/gather?</p>

smth: <p>it would. if you’re worried about that you can put your DataParallel around your model + loss function.<br>
But depending on how many parameters you have in your fully connected layer, it might not work out in terms of speed.</p>

dragen: <p>Hi, If i wrap the loss within forward, let’s say I run on 2 gpus, the forward function will return with [loss1, loss2], and should I sum over the loss1 + loss2 and then backward?<br>
or loss1.backward(), loss2.backward()?</p>

smth: <p>the forward function will receive <code>output = torch.cat([loss1, loss2])</code>, so you can do <code>output.backward(torch.ones(2))</code></p>

GriffinLiang: <p>If output is [loss1, loss2], can I get the final loss as output.sum() ? And then do loss.backward().</p>

henrique: <p>Why wouldn’t you just get the mean like most loss functions do on regular batches?<br>
i.e.</p>
<blockquote>
<p>loss = criterion(output, target).mean()<br>
loss.backward()<br>
seems to work fine</p>
</blockquote>

halahup: <p>What does this do, could you please elaborate?</p>

smth: <p>we are passing a gradient of ones to the backward. Usually, if it’s a scalar output <code>loss</code>, and you do <code>loss.backward()</code>, it’s implied that it’s <code>loss.backward(torch.ones(1))</code>. Because, in this case the loss is actually two elements, <code>output.backward()</code> will give an error asking for gradients.</p>

Dexter_JU: <p>I would have the same question as the guys before,<br>
Can I get the final loss by output.sum()  and then do loss.backward()? (I saw some blog posts doing that way.)<br>
Is that different from what you suggested here?</p>

sunshineatnoon: <p>Same question here, did you find any difference between using sum() and loss.backward(torch.ones(2))?</p>