Semantic segmentation: CrossEntropyLoss with single channel index mask vs. BCELoss with multi-channel one hot vector mask?

username: Yuerno
cooked: <p>Hey all! I’m working with a dataset for segmentation that has your typical RGB images as well as index-based masks (so the pixel values in the mask represent the classes at those positions), and I’m wondering what the best way to train this dataset would be with a simple U-Net. Would it be better to do standard CrossEntropy loss using the index masks, or would it be better to split the masks up into one hot vector format and use BCELoss with each channel? Or is there a better loss for this type of dataset entirely? I’ve seen dice loss, but no efficient implementations for multi-class scenarios.</p>

username: KFrank
cooked: <p>Hello Yuerno -</p>
<aside class="quote no-group" data-username="Yuerno" data-post="1" data-topic="47616" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/y/ccd318/48.png" class="avatar"> Yuerno:</div>
<blockquote>
<p>Hey all! I’m working with a dataset for segmentation that has your typical RGB images as well as index-based masks (so the pixel values in the mask represent the classes at those positions), and I’m wondering what the best way to train this dataset would be with a simple U-Net. Would it be better to do standard CrossEntropy loss using the index masks, or would it be better to split the masks up into one hot vector format and use BCELoss with each channel? Or is there a better loss for this type of dataset entirely? I’ve seen dice loss, but no efficient implementations for multi-class scenarios.</p>
</blockquote>
</aside>
<p>Conventional wisdom and common practice suggest that you<br>
should use CrossEntropyLoss for a multi-class classification<br>
problem.  Reading between the lines, I understand your<br>
segmentation problem to be a straightforward (multi-label)<br>
classification problem, so I would think that CrossEntropyLoss<br>
would be your preferred choice.</p>
<p>Also, if you propose using BCELoss, to be concrete, you need<br>
to specify how you intend to combine the BCELoss for each<br>
of the classes together to get a single loss function to optimize.</p>
<p>Lastly, people argue that dice loss works well (better than<br>
CrossEntropyLoss?) when your classes are imbalanced, but<br>
dice loss – its gradients, etc. – are mathematically unwieldy,<br>
and this is a practical disadvantage.</p>
<p>I would recommend that you start with CrossEntropyLoss,<br>
and only if that isn’t working well on your problem (or maybe<br>
if your classes are highly imbalanced) also try dice loss and<br>
compare its performance to that which you achieved with<br>
CrossEntropyLoss.  (I don’t think you should use just dice loss<br>
without comparing it to CrossEntropyLoss.)</p>
<p>(I’m not aware of any scheme for combining together multiple<br>
BCELosses that people prefer to CrossEntropyLoss.)</p>
<p>Good luck.</p>
<p>K. Frank</p>

username: Yuerno
cooked: <p>Hi Frank! Thanks so much for this detailed response, it really helped clarify things for me. Your assessment in reading between the lines was absolutely correct, so I think I’ll be sticking with just regular CrossEntropyLoss for now and see how that goes.</p>