Loss function for segmentation models

Hi, I’m trying to build a segmentation model with 3 classes.
This is my way of action:
1.my output from the model is :(1,3,512,512)
2. softmax on channel dimenssion.
3. argmax on channel dimension.
4. getting (1,512,512) tensor, correct so far.

When I’m trying to enter to NLLLoss2d , I’m getting an error:
expected … (1,512) , and not (1,512,512)

So I moved to NLLLoss , and tried to flatten my tensors into : (1,262144)
This also didn’t work, I’m getting:
"RuntimeError: multi-target not supported at c:\new-builder_2\win-wheel\pytorch\aten\src\thnn\generic/ClassNLLCriterion.c:21 "

What I’m missing? in Keras it’s so simple.

Please help me understand how to train segmentation models.

You could apply nn.LogSoftmax on your output and nn.NLLLoss for your semgmentation task as you would do it on a simple classification task.
nn.NLLLoss will take multi-dimensional input, so you can just calculate your loss as:

output = model(input)  # raw logit output in shape [1, 3, 512, 512]
loss = criterion(F.log_softmax(output,1), target)  # target in shape [1, 512, 512]

As I mentioned, it didn’t work for me, I’m getting :
“ValueError: Expected target size (1, 512), got torch.Size([1, 512, 512])”

Another attempt I tried, I did exactly as I did in keras and did “to_categorial” like operator, now my target (mask) is (1,3,512,512).

so this is what I did:

outputs = model(inputs)
loss = criterion(F.log_softmax(outputs,1), masks.type(torch.LongTensor))

but I got this error:
RuntimeError: invalid argument 3: only batches of spatial targets supported (3D tensors) but got targets of dimension: 4 at c:\new-builder_2\win-wheel\pytorch\aten\src\thnn\generic/SpatialClassNLLCriterion.c:59

Note: the size of the masks is :(1,3,512,512), and the outputs size is also:(1,3,512,512)

Which PyTorch version are you using?
In older versions you had to use nn.NLLLoss2d as far as I remember. Since 0.3 or 0.4 nn.NLLLoss accepts multi-dimensional inputs.

My Pytorch is 0.4.1…

It’s sucks, because I can go back to Keras, but I still want to understand how to do this…

In that case nn.NLLLoss should work fine.
Sorry for missing the shape information, but your mask shouldn’t have the channel dimension, but store the class indices instead.
Here is a small example code:

batch_size = 1    
nb_classes = 3
h, w = 24, 24
output = torch.randn(batch_size, nb_classes, h, w)
target = torch.empty(batch_size, h, w, dtype=torch.long).random_(nb_classes)
criterion = nn.NLLLoss()
loss = criterion(output, target)