Hello
I want to implement the U-net with Pytorch but I have several questions.
The first question is: The input of U-net is a slice of the abdominal CT scan and The organs must be segmented at the output, should I set the number of U-net output classes to the number of organs?
every slice has a size like 1x512x512
I set U-net input size to one, and output size(number of classes) to number of organs. I mean, is it right?
Second question: how many dimensional should Grand truth be? Suppose we have 10 classes(10 organs). Should I create a tensor, which, in every dimension, is one of the organs maintained by Grand truth(10x512x512)? Or do I create a tensor with one dimension and identify each organ with an independent value(1x512x512)?
If you are using nn.CrossEntropyLoss or nn.NLLLoss as your criterion, the output of your model should have the shape [batch_size, nb_classes, height, width], while your target should be a LongTensor in the shape [batch_size, height, width] containing the class indices.
If you are not dealing with a background class, the number of classes should correspond to the number of organs.
Thank you so much for replying, and another question came to me that is: if I change criterion from nn.CrossEntropyLoss to Dice( or other criterion functions) the shape of the target is important?
Now I am using nn.CrossEntropyLoss but the output after 500 epochs is not satisfied me and it’s not like the output of papers with same situations, so I thought maybe my code is wrong, so I doubted the size of the target and model. So the size of the Target and the model can cause the problem of not learning?
Which size did you double? The spatial size or the number of channels?
If you are using another loss function, e.g. dice loss, you have to make sure to pass the output and target in a specific shape. Have a look at this post to see an implementation of multi-class dice loss.