torch.nn.DataParallel 如何更改输出大小？

Question

看起来使用torch.nn.DataParallel改变了输出大小。尽管在官方文档https://pytorch.org/docs/stable/nn.html#torch.nn.DataParallel中 ，有关大小更改的所有信息如下：

当模块在forward()中返回一个标量（即0维张量）时，该包装器将返回一个长度等于数据并行中使用的设备数量的向量，其中包含每个设备的结果。

我的模块返回 10 个坐标的张量，并且我有 2 个 GPU，我想在其中运行代码。我的 CNN 的最后一层是nn.Linear(500, 10).

import torch
import torch.nn as nn

net = LeNet()    #CNN-class written above
device = torch.device("cuda:0")
net.to(device)
net = nn.DataParallel(net)

#skipped some code, where inputs and targets are loaded from files

output = net(input)
criterion = nn.SmoothL1Loss()
loss = criterion(output, target)

请注意，不调用DataParallel这段代码也可以正常工作。DataParallel当尝试计算损失时会发生运行时错误。

RuntimeError: The size of tensor a (20) must match the size of tensor b (10) at non-singleton dimension 0

如前所述，似乎每个 GPU 的输出大小分别为 10，但随后两个输出合并在一起，这就是大小 20 的来源。

当将 CNN 类中的输出大小从 10 更改为 5 时，它再次开始工作，但我不确定这是正确的解决方案并且 CNN 是否会正常工作。

Answer 1

解决此问题的最简单方法是在向后执行之前对所有 GPU 的损失进行平均。当不同的 GPU 具有不同数量的样本时，这可能会出现问题，但适用于您的情况：

loss = criterion(output, target)
loss = loss.mean()
loss.backward()