Problem with Unet multiclass image segmentation and overlapping classes

fakealex: <p>Hi,<br>
I’m working on a personal project, and I’m trying to use UNet to mask the lead and the big boxes from the ECG paper into two different classes. I created my dataset containing 200 images of ECG paper, lead masks, and big boxes masks. However, my problem is that the model always predicts both classes the same way, with both classes being visible in the other class’s prediction. for example:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/7/7/77bc6499d8368e39287fe54e629088b9eff9bb95.png" alt="Figure_1" data-base62-sha1="h5erQcS5RL6pmE9BSY47m04j0kl" width="640" height="480"><br>
my image shape is [N, C, H, W] (batch, class, height, width) and the class dims are:<br>
[0, 0] if the pixel is not in any classes, example: [N, [0, 0], H, W]<br>
[1, 0] if the pixel is in the class lead, example: [N, [1, 0], H, W]<br>
[0, 1] if the pixel is in the class boxes, example: [N, [0, 1], H, W]<br>
[1, 1] if the pixel is in both classes, example: [N, [1, 1], H, W]<br>
I’m using Dice Loss and the Adam Optimizer. The model outputs in the shape of [N, 2, H, W] (2 for the two classes), and I squeeze then split them using numpy: [0, :, :] for the lead class and [1, :, :] for the boxes class into their respective variables.<br>
I don’t understand why the model predicts both classes in both classes. Could it be that my way of dividing them into two different arrays?<br>
(there are some minor differences between both class predictions, so I know they are not exact copy of each other.)</p>

ptrblck: <p>Could you check how your criterion processes the model outputs and if multi-label targets are supported? For the sake of debugging you could also pass the model output to <code>nn.BCEWithLogitsLoss</code> and compare it to your current approach.</p>

fakealex: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="207472">
<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>criterion</p>
</blockquote>
</aside>
<p>Hi,<br>
the Dice Loss is the only part that I didn’t write myself and I found its implementation online:</p>
<pre><code class="lang-auto">class MultilabelDiceLoss(nn.Module):
    def __init__(self, smooth=1.0):
        super(MultilabelDiceLoss, self).__init__()
        self.smooth = smooth

    def forward(self, inputs, targets):
        inputs = torch.tanh(inputs)
        targets = targets.float()

        inputs = inputs.view(inputs.size(0), inputs.size(1), -1)
        targets = targets.view(targets.size(0), targets.size(1), -1)
        
        inputs = (inputs + 1) / 2
        targets = (targets + 1) / 2

        intersection = (inputs * targets).sum(2)
        union = inputs.sum(2) + targets.sum(2)
        
        dice = (2 * intersection + self.smooth) / (union + self.smooth)
        
        loss = 1 - dice
        return loss.mean()
</code></pre>
<p>and about the <code>nn.BCEWithLogitsLoss</code>, it gives me negative loss values from the first epoch.<br>
also, one thing that I forgot to mention I that the pixel values are normalized between -1 and 1. so that might be the reason for negative loss values with <code>nn.BCEWithLogitsLoss</code>.</p>

fakealex: <p>I changed the pixel values to be between 0 and 1, which fixed the negative loss value. then I tried to train it with <code>nn.BCEWithLogitsLoss</code>. after 60 epochs, the results were the same as with Dice loss.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/a/da56668c92e72c7d6f1ef62a494df9d5648458c2.png" data-download-href="https://discuss.pytorch.org/uploads/default/da56668c92e72c7d6f1ef62a494df9d5648458c2" title="Figure_2"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/a/da56668c92e72c7d6f1ef62a494df9d5648458c2_2_690x334.png" alt="Figure_2" data-base62-sha1="v9vcsswmLt5Qrmj8l1bCfuFhJoS" width="690" height="334" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/a/da56668c92e72c7d6f1ef62a494df9d5648458c2_2_690x334.png, https://discuss.pytorch.org/uploads/default/optimized/3X/d/a/da56668c92e72c7d6f1ef62a494df9d5648458c2_2_1035x501.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/d/a/da56668c92e72c7d6f1ef62a494df9d5648458c2.png 2x" data-dominant-color="CCBBD1"><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">Figure_2</span><span class="informations">1366×663 50.7 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>