parallel-processing - torch.nn.DataParallel 如何改变输出大小？

Question

看起来像使用 torch.nn.DataParallel改变输出大小。
虽然在官方文档 https://pytorch.org/docs/stable/nn.html#torch.nn.DataParallel
有关尺寸变化的所有信息如下:

When module returns a scalar (i.e., 0-dimensional tensor) in forward(), this wrapper will return a vector of length equal to number of devices used in data parallelism, containing the result from each device.

我的模块返回 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

最简单 solution解决这个问题的方法是在向后做之前取所有 GPU 的损失平均值。当不同的 GPU 具有不同数量的样本时，这可能会出现问题，但适用于您的情况:

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