Loss function for 3D multi-class segmentation

Hello,

I’m trying to experiment with 3D multi-class segmentation, very much like the lesson 14, except it’s multi-class and 3D ;), but i’m struggling to understand how to make a loss function for it.

the output’s shape of my network is (batch, nb_of_classes, dimx, dimy, dimz) and my ground truth’s shape is (batch, 1, dimx, dimy, dimz)

when calling the following code:

learn = ConvLearner.from_model_data(net, dataset)
learn.fit(0.01, 1)

i’m getting the following error

ValueError: Expected target size (10, 9, 9, 9), got torch.Size([10, 2, 9, 9, 9])

(reached from ‘self.crit(preds, y)’ from fastai/model.py, which is using nll_loss by default

from there, I’m not quite sure how to implement a proper loss function which would work for that 3D ground truth.

Googling a bit, I’ve seen 3D segmentation code implemented with Keras (and it is just providing a loss function as ‘categorical_crossentropy’), but I didn’t find any example for pytorch.

Could someone help me?

Answering my own question.

I found my bug, i should read the documentation more carefully…
the ground truth’s shape should be (batch, dimx, dimy, dimz)

I got confused with the fact that the input should be unqueezed.

What was the loss function that was used here ?