UNET model not learning

username: Tanvir_Roshid
cooked: <p>I have a UNET model loaded from the SMP library to semantically segment the DDSM database:</p>
<pre><code class="lang-auto">class LesionSegmentation(nn.Module):
    def __init__(self):
        super(LesionSegmentation, self).__init__()

        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model = smp.Unet(
            encoder_name="resnet34",
            encoder_weights="imagenet",
            in_channels=1,
            classes=1,
            activation='sigmoid'
        )

    def forward(self, x):
        # Perform forward pass through model
        mask = self.model(x)
        return mask
</code></pre>
<p>I am using the Dice Loss:</p>
<pre><code class="lang-auto">def dice_score(y_true, y_pred):
    return (2 * (y_true * y_pred).sum() + 1e-15) / (y_true.sum() + y_pred.sum() + 1e-15)
</code></pre>
<p>Images and masks are both B, 1, 224, 224</p>
<p>So far the Dice Loss has proven to be the best in terms of accurately representing my model because other losses, e.g., BCE will be really low (0.005) with a mean IOU of 0.1</p>
<p>My problem is, no matter what I edit, e.g., loss fun, encoder, model etc. My results are the same:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/9/9953c811610e63e4867e84eb354a6e1961ec08bb.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/9953c811610e63e4867e84eb354a6e1961ec08bb" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/9/9953c811610e63e4867e84eb354a6e1961ec08bb_2_690x400.jpeg" alt="image" data-base62-sha1="lSoylKWFhmF77defMIWBpdGnqKD" width="690" height="400" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/9/9953c811610e63e4867e84eb354a6e1961ec08bb_2_690x400.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/9/9/9953c811610e63e4867e84eb354a6e1961ec08bb_2_1035x600.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/9/9/9953c811610e63e4867e84eb354a6e1961ec08bb_2_1380x800.jpeg 2x" data-dominant-color="F9F8F7"><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">image</span><span class="informations">1579×916 141 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>The problem is, I don’t know where to go from here. I’ve trialled all sorts of loss functions because I thought that maybe there is a class imbalance for my dataset in terms of lesion/nonlesion ratio</p>
<p>This is the dataset I’m using: <a href="https://www.kaggle.com/datasets/cheddad/miniddsm2?resource=download-directory" class="inline-onebox" rel="noopener nofollow ugc">The Complete Mini-DDSM | Kaggle</a></p>

username: tom
cooked: <p>I don’t know the dataset, but your hunch sounds very plausible. In my experience, it can be tricky to compensate class imbalance through the loss function. In our book we use cropping to reduce the imbalance in the sampling (a long time ago, <a href="https://discuss.pytorch.org/t/calibrating-probability-output-and-threshold/128303/2">someone had a similar issue</a>).</p>
<p>Best regards</p>
<p>Thomas</p>

username: Tanvir_Roshid
cooked: <p>Thank you <a class="mention" href="/u/tom">@tom</a></p>
<p>I’ll close and take a look</p>
<p>See if i can manage it on my own</p>

username: Tanvir_Roshid
cooked: <p>Hi <a class="mention" href="/u/tom">@tom</a></p>
<p>Wanted to officially confirm that this was the correct answer. I used a bounding box on the mammograms and cropped the image and mask using this technique.</p>
<p>The training scores boosted massively from the first epoch.</p>
<p>Albeit, the validation scores are terrible. Likely because the training and validation don’t represent each other, so I’ll try finding a <strong>sweet</strong> spot between:</p>
<ul>
<li>Representing a mammogram</li>
<li>Balancing the class imbalance</li>
</ul>
<p><img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji only-emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>