nn.CrossEntropyLoss用于图像分割

当你想将nn.CrossEntropyLoss用于二分类及多分类图像分割时,应该如下设置:

你的分割模型的输出output应该是一个 大小为[batch_size, nb_classes, height ,width]的tensor。
对应的掩码标签target应该是一个大小为[batch_size, height, width] 的LongTensor,并且每一个像素的值取值范围是[0, nb_classes-1] (有些数据集中将掩码中的无用像素设为255,在处理时,可以不用处理它,保留255。只需要将参数ignore_index设为255即可)。
注:以上的nb_classes中包含背景类

示例

from torch import nn

loss_fun = nn.CrossEntropyLoss(ignore_index=255)
loss = loss_fun(output, target.long())

参考:CrossEntropyLoss用于图像分割