看起来使用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)
Run Code Online (Sandbox Code Playgroud)
请注意,不调用DataParallel这段代码也可以正常工作。DataParallel当尝试计算损失时会发生运行时错误。
RuntimeError: The size of tensor a (20) must match the size of tensor b (10) at non-singleton …Run Code Online (Sandbox Code Playgroud)