使用pytorch进行多类图像分割中的ground truth形状

Posted 2023-03-27

【中文标题】使用pytorch进行多类图像分割中的ground truth形状

【英文标题】：Shape of ground truth in multiclass image segmentation with pytorch

【发布时间】：2021-05-18 13:17:33

【问题描述】：

我正在处理 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]

我该如何解决这个问题？

【问题讨论】：

请添加U-Net模型的超链接

模型期望的形状是什么？在引发此错误之前，可能有一条 assert 行失败。

@Abhi25t 嗨，这是我使用的模型的链接：github.com/milesial/Pytorch-UNet.git。

模型的预期输出形状为 [2,5,128,128]。我找不到对应的断言语句

【参考方案1】：

您似乎正在使用nn.CrossEntropyLoss 或nn.functional.cross_entropy

我也遇到了同样的错误。

CrossEntropyLoss 通常用于分类用例。

如果您的目标是归一化张量，其值为[0, 1]，您可以使用nn.BCELoss 或nn.functional.binary_cross_entropy_with_logits。这对我来说很有效，因为我们为每个类使用单独的掩码 - 它变成了一个二元交叉熵问题。