Pytorch语义分割损失函数

Question

我是新手，对分割模型不太了解。 我想使用 deeplabv3_resnet50 模型。 我的图像尺寸为 (256, 256, 3) ，标签尺寸为 (256, 256)。标签中的每个像素都有一个类别值（0-4）。DataLoader 中设置的批量大小为 32。 因此，我的输入批量形状为 [32, 3, 256, 256]，相应目标的形状为 [32, 256, 256]。我认为这是正确的。

我试图使用 nn.BCEWithLogitsLoss()。

1. 这是我这种情况下正确的损失函数吗？还是应该使用 CrossEntropy？
2. 如果这是正确的，那么模型的输出是 [32, 5, 256, 256]。每个图像预测的形状为 [5,256, 256]，层 0 是否表示类别 0 的未归一化概率？为了使 [32, 256, 256] 张量与目标匹配并输入到 BCEWithLogitsLoss 中，是否需要将未归一化的概率转换为类别？
3. 如果我应该使用 CrossEntropy，那么输出和标签的大小应该是多少？

谢谢大家。

- KKKcat

Answer 1

你使用了错误的损失函数。 nn.BCEWithLogitsLoss() 代表二元交叉熵损失函数：用于二元标签。在你的情况下，你有5个标签（0..4）。

你应该使用 nn.CrossEntropyLoss：一种为离散标签设计的损失函数，适用于非二元情况。

你的模型应输出一个形状为 [32, 5, 256, 256] 的张量：对于批处理中的32个图像中的每个像素，它应输出一个5维向量的logits。logits是每个类别的原始分数，稍后将使用softmax函数归一化为类概率。

为了数值稳定和计算效率，nn.CrossEntropyLoss 不需要你显式地计算logits的softmax，而是在内部为你执行此操作。正如文档所述：“该标准将LogSoftmax和NLLLoss结合在一个单一的类中。”

- Shai

---

明白。如果我之后想要计算IOU或像素准确度，那么我应该将输出设置为[32, 256, 256]（也许使用output.argmax(dim=1)），以匹配我的标签吗？ - KKKcat

@KKKcat 对通道维度执行 argmax 操作即可得到预测标签。 - Shai

Answer 2

假设你正在处理5个类别，那么你应该使用CrossEntropyLoss。二元交叉熵，顾名思义，是一种损失函数，用于在具有二进制分割映射的情况下使用。

在PyTorch中，CrossEntropy函数期望从模型输出的形状为[batch, num_classes, H, W]（直接传递给损失函数），而ground truth的形状为[batch, H, W]，其中H，W在您的情况下为256，256。此外，请确保通过对张量调用.long()将ground truth的类型设置为long

- Atharva Dubey