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

时间:2019-08-22 10:50:25

标签: parallel-processing conv-neural-network pytorch

似乎使用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可以正常工作。

1 个答案:

答案 0 :(得分:0)

此问题最简单的solution是在向后进行操作之前,先考虑所有GPU的损失。当不同的GPU具有不同的样本数量时,这可能会出现问题,但将适合您的情况:

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