Can I use one-hot encoded output for segmentation in Pytorch, with focal and dice losses?

Question

know that for classification using a neural network and CrossEntropy Loss, we need one-hot encoded output, but in PyTorch, the CrossEntropy loss does not accept one-hot encoded targets, and we should give it the labels, directly and in the normal format.

Now, I am wondering if this is the same for image segmentation tasks, where the loss function is the dice loss or focal loss, etc. i.e. Is it ok if I one-hot encode the target mask for segmentation similar to TensorFlow, or I cannot do that similar to classification task in Pytorch?

Answer 1

For classification, it seems that in the latest version of PyTorch, cross-entropy accepts one_hot encoded labels as well.

For segmentation, PyTorch does not have a Dice loss implementation, hence it can be implemented in any way.