How many output channels in U-net for multiclass segmentation?

Looking at some implementations of Unet, I believe the output segmentation mask will have n_classes number of channels. So n_classes, h,w. I tried looking at the fastai source code but couldn’t figure out the size of that. Any idea what size it is? And why not just have 1 channel, and all the classes in that one channel, instead of having one class per channel?

The default height/width should be that of your input image.

The reason for multiple channels is the model doesn’t just predict a single output at each location. For each pixel the model predicts log probabilities for all classes, the same as you would for a normal multi-class classification problem.


Oh okay, so essentially the output is one-hot encoded so number of output channels is the number of classes.

In this case of multi-class, would it be okay to use binary cross entropy? or will cross entropy for also work?

1 Like

Does someone know a way to onehot encode the label in multiclass form ? To go from Bx1xHxW to BxCxHxW easily ?