问题描述
我试图训练 UNET,输入大小为 [3,128,128],相应的掩码是 [1,128],它直接包含类(而不是像素,它将包含类号 - 1,2)。我正在尝试解决一个两类问题,因此我的掩码包含 1,2 作为标签。现在我将我的图像发送到模型,预测掩码的维度是 [2,128]。现在要训练模型,我选择 16 作为批量大小。所以,现在我输入为 [16,3,128] 所以预测维度是 [16,2,128]。但是我有真实的掩码 [16,1,128]。现在如何在 Pytorch 中应用交叉熵损失?我已尝试如下并收到以下错误。能否请你帮忙?提前致谢。
lr = 0.1 # 0.1
criterion = nn.CrossEntropyLoss() #nn.L1Loss()
optimizer = optim.SGD(model.parameters(),lr=lr,momentum=0.9,nesterov=True,weight_decay=0.0001)
is_train = True
is_pretrain = False
acc_best = 0
total_epoch = 30
if is_train is True:
# Training
for epoch in range(total_epoch):
model.train()
tims = time.time()
for i,(images,labels) in enumerate(train_loader):
images = Variable(images.permute(0,2).float().cuda())
labels = Variable(labels.type(torch.LongTensor).cuda())
# Forward + Backward + Optimize
optimizer.zero_grad()
outputs = model(images)
outputs = outputs.type(torch.float)
print('predictedLabelsType:',outputs[0].type())
print('ActualLabelsType:',labels[0].type())
print('shape of predicted outputs:',outputs.shape)
print('shape of groundtruth masks:',labels.shape)
loss = criterion(outputs,labels)
loss.backward()
optimizer.step()
我的输出如下:
predictedLabelsType: torch.cuda.FloatTensor
ActualLabelsType: torch.cuda.LongTensor
shape of predicted outputs: torch.Size([16,128])
shape of groundtruth masks: torch.Size([16,128])
---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
<ipython-input-8-b692a8d536a9> in <module>()
52 print('shape of predicted outputs:',outputs.shape)
53 print('shape of groundtruth masks:',labels.shape)
---> 54 loss = criterion(outputs,labels)
55 loss.backward()
56 optimizer.step()
3 frames
/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py in nll_loss(input,target,weight,size_average,ignore_index,reduce,reduction)
2385 ret = torch._C._nn.nll_loss(input,_Reduction.get_enum(reduction),ignore_index)
2386 elif dim == 4:
-> 2387 ret = torch._C._nn.nll_loss2d(input,ignore_index)
2388 else:
2389 # dim == 3 or dim > 4
RuntimeError: 1only batches of spatial targets supported (3D tensors) but got targets of size: : [16,128]
您能否提出错误在哪里以及这个 CrossEntropyLoss() 在 Pytorch 中如何工作以进行图像分割?我在这里缺少什么?我曾尝试将目标大小重塑为 [16,128],这导致了另一个错误。非常感谢!
解决方法
documentation 指定如果输入是形状 (N,C,d1,d2)
,那么目标必须是形状 (N,d2)
。相反,您的目标是形状 (N,1,d2)
,因此您需要删除不必要的单一维度。
loss = criterion(output,labels.squeeze(1))
如果您从此更改中收到另一个错误,则说明您的代码存在另一个问题,但这是 CrossEntropyLoss
的正确张量形状。