Should i use nn.BCEWithLogitsLoss() or Cross Entropy loss for segmentation

username: ayadav01
cooked: <p>I am trying to build a simple U-Net for segmentation. The input image as well as the labels has shape (1 x width x height). The pixel values in the label image is either 0 or 1. I am a  beginner to deep learning and just started with pytorch so just want to make sure i am using the right loss function for this task. Will it be better to use binary cross entropy or categorical cross entropy for this task? Also, how can i pass weighted values to the loss function such that it gives more weight to the 1 pixels rather than 0 pixels. Thanks!</p>

username: ptrblck
cooked: <p>For a binary segmentation use case you could use <code>nn.BCEWithLogitsLoss</code> and use the <code>pos_weight</code> argument to weight the positive or negative class more.</p>

username: ayadav01
cooked: <p>ok so if i am using <code>nn.BCEWithLogitsLoss</code>, that means i only have one class output. does that mean i only have to pass only value for the <code>pos_weight</code> argument ? and is there a general formula to find out how much weight to give to each class ?</p>

username: ptrblck
cooked: <p>Yes, you should pass a single value to <code>pos_weight</code>. From the <a href="https://pytorch.org/docs/stable/nn.html?highlight=bcewithlogitsloss#torch.nn.BCEWithLogitsLoss">docs</a>:</p>
<blockquote>
<p>For example, if a dataset contains 100 positive and 300 negative examples of a single class, then pos_weight for the class should be equal to <code>300/100=3 </code>. The loss would act as if the dataset contains <code>3 * 100=300</code> positive examples.</p>
</blockquote>

username: Patrice
cooked: <p>Hello Mr. Ptrblck,<br>
Please, how can I determine the value to pass to pos_weight? In my case, all the images contain both background and foreground (positive and negative). However, the number of foreground pixels (positive) is much smaller than that of background pixels (negative). Please, could anyone help me solve it?<br>
Any suggestions or comments would be highly appreciated.</p>

username: ptrblck
cooked: <p>You could count the number of positive and negative pixels for your complete training dataset and use the average of these counts to calculate the <code>pos_weight</code> or alternatively you could calculate the <code>pos_weigth</code> for the current batch by counting the positive and negative pixels and use the functional API via <code>F.binary_cross_entropy_with_logits</code>.</p>

username: Patrice
cooked: <aside class="quote no-group" data-username="ptrblck" data-post="6" data-topic="79476">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>You could count the number of positive and negative pixels for your complete training dataset and use the average of these counts</p>
</blockquote>
</aside>
<p>Please sir, how can I achieve this?<br>
I have computed the mean() of a given image, which gave the fraction of 1s in the image. Then, I summed up these values to obtain that of the whole training data. I passed the obtained value to pos_weight. but there was no improvement.</p>

username: ptrblck
cooked: <p>You could select a small subset of your dataset (e.g. just 10 samples) and validate your approach (e.g. if the right <code>pos_weight</code> was used or to further fine tune the <code>pos_weight</code>).<br>
Once your model training behaves as desired, you could scale it up again.</p>

username: ayadav01
cooked: <p>Hi Patrice, For a segmentation model, i would suggest you to use a soft dice loss rather than using weighted nn.BCEWithLogitsLoss(). I don’t think you would see much improvments even with a weight cross entropy loss. You can implement soft dice loss like this:</p>
<pre><code class="lang-auto">import torch.nn as nn
import torch

class SoftDiceLoss(nn.Module):
    def __init__(self, weight=None, size_average=True):
        super(SoftDiceLoss, self).__init__()

    def forward(self, logits, targets):
        smooth = 1
        num = targets.size(0)
        """
       I am assuming the model does not have sigmoid layer in the end. if that is the case, change torch.sigmoid(logits) to simply logits
        """
        probs = torch.sigmoid(logits)
        m1 = probs.view(num, -1)
        m2 = targets.view(num, -1)
        intersection = (m1 * m2)

        score = 2. * (intersection.sum(1) + smooth) / (m1.sum(1) + m2.sum(1) + smooth)
        score = 1 - score.sum() / num
        return score

criterion = SoftDiceLoss().to(device) # to put on cuda or cpu
# In the training loop
output = model(images)
loss = criterion(output, gt_masks)
# Then do the backwards and optimizer step
</code></pre>

username: Patrice
cooked: <aside class="quote no-group" data-username="ayadav01" data-post="9" data-topic="79476">
<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/a/c67d28/48.png" class="avatar"> ayadav01:</div>
<blockquote>
<p>Hi Patrice, For a segmentation model, i would suggest you to use a soft dice loss rather than using weighted nn.BCEWithLogitsLoss()</p>
</blockquote>
</aside>
<p>Hello Anil, thank you very much  for your suggestion. I will give it a try after fixing my computer. I am facing some hardware problems presently.</p>

username: Patrice
cooked: <aside class="quote no-group" data-username="ptrblck" data-post="8" data-topic="79476">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>You could select a small subset of your dataset (e.g. just 10 samples) and validate your approach (e.g. if the right <code>pos_weight</code> was used or to further fine tune the <code>pos_weight</code> ).</p>
</blockquote>
</aside>
<p>Ok sir, I will try that.</p>

username: morisa66
cooked: <p>I recommend the loss function of BCELogisticLoss <img src="https://discuss.pytorch.org/images/emoji/apple/grin.png?v=9" title=":grin:" class="emoji" alt=":grin:"></p>