UNet Multiclass Loss Function Selection

droid: <p>I have a vanilla implementation of UNet, which I want to use for multiclass segmentation (where each pixel can belong to many classes). I am interested in advice on which loss function to select in this application.</p>
<p><a href="https://discuss.pytorch.org/t/multi-class-cross-entropy-loss-function-implementation-in-pytorch/19077/83">This long thread</a> suggests using <code>CrossEntropyLoss</code> at first, before recommending <code>BCELoss</code>. Is one of these methods preferred over the other? Is there some other, better method?</p>
<p>If using <code>CrossEntropy</code>, I assume I can pass in an entire batch at once to calculate the loss, while with <code>BCELoss</code>, I assume I must individually find the loss for each output channel, per image, and sum them together to get an overall loss value. If this is true, <code>BCELoss</code> sounds like a lot more work in this scenario.</p>
<p>Any tips are appreciated.</p>

KFrank: <p>Hi Droid!</p>
<aside class="quote no-group" data-username="droid" data-post="1" data-topic="138106" 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/d/54ee81/48.png" class="avatar"> droid:</div>
<blockquote>
<p>I have a vanilla implementation of UNet, which I want to use for multiclass segmentation (where each pixel can belong to many classes).</p>
</blockquote>
</aside>
<p>Let me distinguish between a (single-label) multi-class problem and<br>
a <em>multi-label,</em> multi-class problem.  In both cases you have multiple<br>
classes (one of which might be a “background” or catch-all “other”<br>
class).</p>
<p>In the single-label case, each pixel (or more generally, each item<br>
to be classified) belongs to <em>exactly one</em> class.  <code>CrossEntropyLoss</code><br>
is the natural loss criterion in this case.</p>
<p>In the multi-label case, each pixel belongs to any number of the<br>
classes, including no class or all classes.  <code>BCEWithLogitsLoss</code><br>
is the natural loss criterion to use for a multi-label, multi-class<br>
problem, and it works out of the box for this use case without any<br>
additional fussing around.</p>
<p>You say that “each pixel can belong to many classes,” so I understand<br>
you to be working with a multi-label, multi-class problem.</p>
<p>The multi-label problem should be understood as number-of-classes<br>
binary classification problems all being run in parallel in the same<br>
network.  Is the pixel in class-A, yes or no? (a binary problem);<br>
independently, is the pixel in class-B, yes or no? (a second binary<br>
problem); and so on.</p>
<p>(Please note that <code>BCEWithLogitsLoss</code> is to be preferred to <code>BCELoss</code><br>
for reasons of numerical stability.)</p>
<p>Best.</p>
<p>K. Frank</p>

droid: <p>Thanks KFrank. This makes sense to me. I just needed a nudge in the right direction! Thanks also for clarifying multi-class and multi-label problems.</p>
<p>I have a related followup question. <code>BCEWithLogitsLoss</code> seems be working quite well with my dataset (&gt;99% accuracy). However, I am completely missing one of my output channels (the green channel in the below example). It seems like because this class is small in magnitude, compared to the other classes, the network is not learning this class at all.</p>
<p>Would you recommend scaling the class weights for <code>BCEWithLogitsLoss</code>, or perhaps something else? I do not have a background class – just 3 output channels for R, G, B. I create the predicted output image using <code>torch.round(F.sigmoid(model(data)))</code> to force every pixel to 0 or 1 for each class.</p>
<p>I have not used channel scaling or a background channel before, so I’m not sure how to go about doing that.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/1/310bd71a9b176064aba5a33771ef15efcfb1e879.png" data-download-href="https://discuss.pytorch.org/uploads/default/310bd71a9b176064aba5a33771ef15efcfb1e879" title="example"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/1/310bd71a9b176064aba5a33771ef15efcfb1e879_2_517x198.png" alt="example" data-base62-sha1="6ZSLlxX6FucKTGbv6RXJATKYZst" width="517" height="198" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/1/310bd71a9b176064aba5a33771ef15efcfb1e879_2_517x198.png, https://discuss.pytorch.org/uploads/default/optimized/3X/3/1/310bd71a9b176064aba5a33771ef15efcfb1e879_2_775x297.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/3/1/310bd71a9b176064aba5a33771ef15efcfb1e879.png 2x" data-dominant-color="4D4198"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">example</span><span class="informations">957×368 243 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

KFrank: <p>Hi Droid!</p>
<aside class="quote no-group" data-username="droid" data-post="3" data-topic="138106" 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/d/54ee81/48.png" class="avatar"> droid:</div>
<blockquote>
<p>I do not have a background class – just 3 output channels for R, G, B.</p>
</blockquote>
</aside>
<p>First, a clarifying question about your multi-class segmentation<br>
problem:</p>
<p>In the “Mask” (and “Predicted”) image you post I see four colors:<br>
black, green, blue, and magenta.  You might have other images<br>
with more colors, but it looks to me like you might be performing<br>
single-label, four-class (or perhaps a few more than four-class)<br>
segmentation.</p>
<p>I understand that magenta is made of red and blue, so you could<br>
interpret it as being in the both the “red” class and the “blue” class,<br>
but is there any substantive sense in which the pixels making up<br>
the QR code in your “Image” image are logically both in a “red”<br>
and “blue” class, or did you just happen to label the QR pixels with<br>
magenta in your mask?</p>
<p>Just because you choose to label pixels in your mask with saturated<br>
RGB colors doesn’t mean that your problem is naturally a multi-label,<br>
three-class (RGB) problem.  Is it possible that single-label, multi-class<br>
segmentation is a better logical fit for your use case?</p>
<aside class="quote no-group">
<blockquote>
<p>I am completely missing one of my output channels (the green channel in the below example). It seems like because this class is small in magnitude, compared to the other classes, the network is not learning this class at all.</p>
</blockquote>
</aside>
<p>Whether you treat this as single-label or multi-label, the fact that one<br>
of your classes (green) is significantly less common than the others<br>
can lead to an otherwise well-trained network not predicting the “rare”<br>
class at all.</p>
<p>Class weights, where you weight the green class more heavily, make<br>
sense in such a case, and both <code>CrossEntropyLoss</code> (for the single-label<br>
case) and <code>BCEWithLogitsLoss</code> (for the multi-label case) support class<br>
weights.  A common, sensible approach is to weight each class in<br>
inverse proportion to how often it occurs in your training set.</p>
<aside class="quote no-group">
<blockquote>
<p>I create the predicted output image using <code>torch.round(F.sigmoid(model(data)))</code> to force every pixel to 0 or 1 for each class.</p>
</blockquote>
</aside>
<p>As a side note, you don’t need to use <code>sigmoid()</code> / <code>round()</code> to convert<br>
your floating-point network output to a binary prediction.  Thresholding<br>
against zero, <code>data &gt; 0.0</code>, is equivalent, and I prefer it stylistically.</p>
<p>Best.</p>
<p>K. Frank</p>

droid: <p>You are correct that there are only 4 channels (RGB + background as black) in my masks. Each channel/class is binary with the black background as zeros. I have my model setup for only 3 output channels.</p>
<p>I do want to keep the multi-label aspect of this classification so that I can analyze each channel independently in a postprocessing step. So it is important to me that the network learns “class independence”, so to speak (although that might be the wrong term). For instance, I want the network to learn there are “blue plate pixels” under the QR code, so that I can analyze the blue channel without worrying about it having a QR code-shaped hole through it. I want to be able to measure the blue plate characteristics regardless of the QR code or any other overlapping label.</p>
<p>The same logic follows for the other channels, where it’s important to have overlapping segments for this application. I hope that makes sense.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/1/413e8639c1d3533a86824ee0c5e8155e2e76221a.png" data-download-href="https://discuss.pytorch.org/uploads/default/413e8639c1d3533a86824ee0c5e8155e2e76221a" title="Capture"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/1/413e8639c1d3533a86824ee0c5e8155e2e76221a_2_289x250.png" alt="Capture" data-base62-sha1="9jaZAiDC0XPhcsBAGpZpI46hVqi" width="289" height="250" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/1/413e8639c1d3533a86824ee0c5e8155e2e76221a_2_289x250.png, https://discuss.pytorch.org/uploads/default/optimized/3X/4/1/413e8639c1d3533a86824ee0c5e8155e2e76221a_2_433x375.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/4/1/413e8639c1d3533a86824ee0c5e8155e2e76221a_2_578x500.png 2x" data-dominant-color="00098F"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Capture</span><span class="informations">611×527 33.5 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>With this in mind, I believe <code>BCEWithLogitsLoss</code> is still the way to go based on your earlier description. However, I’m having some trouble setting up the weights with this loss function (see below). I found <a href="https://discuss.pytorch.org/t/weight-vs-pos-weight-in-nn-bcewithlogitsloss/114859/4">this post</a> (from you) that offers some explanation but I’m still confused how to get the weights up and running.</p>
<pre><code class="lang-auto"># red, blue, green channel weights = 6, 1, 30
loss_func = nn.BCEWithLogitsLoss(
    pos_weight=torch.Tensor([1, 30, 6]).repeat(BATCH_SIZE, 1))

# RuntimeError: The size of tensor a (3) must match the size of
# tensor b (256) at non-singleton dimension 3
</code></pre>
<p>Any advice?</p>

KFrank: <p>Hi Droid!</p>
<aside class="quote no-group" data-username="droid" data-post="5" data-topic="138106" 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/d/54ee81/48.png" class="avatar"> droid:</div>
<blockquote>
<p>The same logic follows for the other channels, where it’s important to have overlapping segments for this application.</p>
</blockquote>
</aside>
<p>Yes, it sounds like you do have a multi-label, three-class problem.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto"># red, blue, green channel weights = 6, 1, 30
loss_func = nn.BCEWithLogitsLoss(
    pos_weight=torch.Tensor([1, 30, 6]).repeat(BATCH_SIZE, 1))

# RuntimeError: The size of tensor a (3) must match the size of
# tensor b (256) at non-singleton dimension 3
</code></pre>
</blockquote>
</aside>
<p>Unfortunately, the documentation for <a href="https://pytorch.org/docs/1.10.0/generated/torch.nn.BCEWithLogitsLoss.html" rel="noopener nofollow ugc">BCEWithLogitsLoss</a> is pretty<br>
much silent about your use case (sometimes called the <em>K dimensional<br>
case</em> in other parts of the pytorch documentation).  Specifically, your<br>
prediction (and mask) don’t have shape <code>[nBatch, nChannel]</code>, but<br>
instead also have spatial dimensions, so that its shape is all of<br>
<code>[nBatch, nChannel, height, width]</code>.  That is you are not just<br>
making a multi-label (<code>nChannel</code>) prediction for each sample in the<br>
batch, but for each pixel in the image individually.</p>
<p>In this case, <code>BCEWithLogitsLoss</code>’s <code>weight</code> and <code>pos_weight</code> tensors<br>
(passed in as constructor arguments to <code>BCEWithLogitsLoss</code>) are<br>
<em>broadcast</em> over <code>pred</code> and <code>mask</code> (called <code>input</code> and <code>target</code> in the<br>
documentation).  Therefore, to make your <code>pos_weight</code> tensor properly<br>
broadcastable, you need to add two singleton dimensions to its end<br>
(that will then be broadcast over the <code>height</code> and <code>width</code> dimensions).<br>
It doesn’t matter how you do this – one way is to use <code>unsqueeze()</code>.</p>
<p>Thus, this should work:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.9.0'
&gt;&gt;&gt; _ = torch.manual_seed (2021)
&gt;&gt;&gt; nBatch = 2
&gt;&gt;&gt; nChannel = 3   # R, G, B
&gt;&gt;&gt; height = 10
&gt;&gt;&gt; width = 12
&gt;&gt;&gt; pred = torch.randn (nBatch, nChannel, height, width)   # logits in [-inf, inf]
&gt;&gt;&gt; mask = torch.rand (nBatch, nChannel, height, width)   # probabilities in [0.0, 1.0]
&gt;&gt;&gt; pos_wt = torch.tensor ([6.0, 30.0, 1.0])   # weights are floats
&gt;&gt;&gt; torch.nn.BCEWithLogitsLoss (pos_weight = pos_wt.unsqueeze (-1).unsqueeze (-1)) (pred, mask)
tensor(5.2394)
</code></pre>
<p>Note, <code>mask</code> should be of floating-point type, not integer.  Your mask<br>
entries can be exactly <code>0.0</code> or <code>1.0</code> if you want, but they can be any<br>
value in between, as well.</p>
<p>Best.</p>
<p>K. Frank</p>

droid: <blockquote>
<p>Unfortunately, the documentation for <a href="https://pytorch.org/docs/1.10.0/generated/torch.nn.BCEWithLogitsLoss.html" rel="noopener nofollow ugc">BCEWithLogitsLoss</a> is pretty<br>
much silent about your use case</p>
</blockquote>
<p>Wow, thank you for explaining. What you described is honestly something I would not have figured out by myself reading the docs… For this application, using the right loss function and loss weighting was the critical piece. I don’t think the model learning would have converged without it, even in this simplified MVP. My segmentation is now looking great with extremely high accuracy after I implemented the weighting correctly. <img src="https://discuss.pytorch.org/images/emoji/apple/beers.png?v=12" title=":beers:" class="emoji" alt=":beers:" loading="lazy" width="20" height="20"></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/e/ee815a4256b9efb8c23c8119dea033c0b59aa2d8.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/ee815a4256b9efb8c23c8119dea033c0b59aa2d8" title="Capture.PNG"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/e/ee815a4256b9efb8c23c8119dea033c0b59aa2d8_2_517x257.jpeg" alt="Capture.PNG" data-base62-sha1="y1UM1u8kt5TWG2Kp8iMI2Jci4fm" width="517" height="257" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/e/ee815a4256b9efb8c23c8119dea033c0b59aa2d8_2_517x257.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/e/e/ee815a4256b9efb8c23c8119dea033c0b59aa2d8_2_775x385.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/e/e/ee815a4256b9efb8c23c8119dea033c0b59aa2d8_2_1034x514.jpeg 2x" data-dominant-color="1B1543"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Capture.PNG</span><span class="informations">1913×953 115 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

KiTo: <p>I came across the same problem and I am not able to solve it … Could you share your Unet implementation?</p>