Pytorch语义分割损失函数

Question

我\xe2\x80\x99m 是分割模型的新手。\n我想使用deeplabv3_resnet50模型。\n我的图像具有 shape (256, 256, 3)，我的标签具有 shape (256, 256)。我的标签中的每个像素都有一个类值（0-4）。DataLoader 中设置的批量大小为 32。\n因此，我输入批量的形状为[32, 3, 256, 256]，对应目标的形状为[32, 256, 256]。我相信这是正确的。

\n

我试图使用nn.BCEWithLogitsLoss().

\n

\n
1. 对于我的情况来说，这是正确的损失函数吗？或者我应该使用\nCrossEntropy代替？
\n
2. 如果这是正确的，那么我的模型的输出是[32, 5, 256, 256]. 每个图像预测的形状为[5,256, 256]，第 0 层是否表示第 0 类的未归一化概率？为了使[32, 256, 256]张量与输入的目标相匹配BCEWithLogitsLoss，我是否需要将未归一化的概率转换为类？
\n
3. 如果我应该使用CrossEntropy，我的输出和标签的大小应该是多少？
\n

\n

谢谢大家。

\n

Answer 1

您使用了错误的损失函数。

nn.BCEWithLogitsLoss()代表二进制交叉熵损失：即二进制标签的损失。就您而言，您有 5 个标签 (0..4)。
您应该使用nn.CrossEntropyLoss：为离散标签设计的损失，超出二进制情况。

您的模型应该输出一个形状为 的张量[32, 5, 256, 256]：对于该批次的 32 个图像中的每个像素，它应该输出一个 5 维的logits向量。logits是每个类别的“原始”分数，稍后使用 softmax 函数标准化为类别概率。
为了数值稳定性和计算效率，nn.CrossEntropyLoss不需要您显式计算 logits 的 softmax，而是在内部为您计算。正如文档所示：

该标准将 LogSoftmax 和 NLLLoss 结合在一个类别中。