问题描述
我正在处理 128 x 128 x 3 的单元格图像，并希望将它们分成 5 个类别，包括背景。我首先将目标图像设为 128 x 128，值在 {0,1,2,3,4} 中。但是我发现我必须将目标ground truth设为5通道图像，并且所有值都是0或1：如果像素在第n通道中为1，则应将其分类为第n类。
但是当我将我的模型运行到我从 GitHub 分叉的 Unet 模型中时，我发现在计算交叉熵损失时出现错误。
我最初将输入中的通道数设置为 3，输出中的类数设置为 5。批大小 = 2
这是我的代码：
for i,(x,y) in batch_iter:
input,target = x.to(self.device),y.to(self.device) # send to device (GPU or cpu)
self.optimizer.zero_grad() # zerograd the parameters
out = self.model(input) # one forward pass
loss = self.criterion(out,target) # calculate loss
loss_value = loss.item()
train_losses.append(loss_value)
loss.backward() # one backward pass
self.optimizer.step() # update the parameters
batch_iter.set_description(f'Training: (loss {loss_value:.4f})') # update progressbar
self.training_loss.append(np.mean(train_losses))
self.learning_rate.append(self.optimizer.param_groups[0]['lr'])
batch_iter.close()
和错误信息
RuntimeError: 1only batches of spatial targets supported (3D tensors) but got targets of size: : [2,5,128,128]
我该如何解决这个问题？
解决方法
您似乎正在使用 nn.CrossEntropyLoss 或 nn.functional.cross_entropy
我也遇到了同样的错误。
CrossEntropyLoss 通常用于分类用例。
如果您的目标是值在 [0,1] 中的归一化张量，则可以使用 nn.BCELoss 或 nn.functional.binary_cross_entropy_with_logits。这在我的情况下有效，因为我们为每个类使用单独的掩码 - 它变成了一个二元交叉熵问题。