UNET Segmentation Model Learns only 1-Class [Colab Shared]

aksg87: <p>Hello all,</p>
<p>I am having a strange issue where my model is only learning a specific class in a data-set (Liver class ~ blue). I experimented more and realized if I remove the Liver class, the model can identify the kidneys very well (yellow and red). Interestingly the kidneys are much brighter in these images!</p>
<p>Does this essentially mean my model is short on parameters for predicting all organs? I would think adding classes would enhance the original model.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/1/0/102c9a1cac3dc38b6333aa44d441531e55683bf8.png" alt="image" data-base62-sha1="2j5cksE3o0ja54IaM8R2zvdwfW0" width="645" height="460"><br>
Image: Source, Ground, Prediction</p>
<p><strong>Link to Original Colab Notebook</strong>. All data downloads into the notebook automatically! (You can just click ‘Open in playground’ to run the code).</p>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://colab.research.google.com/drive/1cDXzCnTaVDEe3Dnhafp-Rdu6rKs3BuXI">
  <header class="source">
      <img src="https://ssl.gstatic.com/colaboratory-static/common/f6c736cbde16632a9fcd16d9ab1970b1/img/favicon.ico" class="site-icon" width="128" height="128">

      <a href="https://colab.research.google.com/drive/1cDXzCnTaVDEe3Dnhafp-Rdu6rKs3BuXI" target="_blank" rel="noopener nofollow ugc">colab.research.google.com</a>
  </header>

  <article class="onebox-body">
    <img width="260" height="260" src="https://colab.research.google.com/img/colab_favicon_256px.png" class="thumbnail onebox-avatar">

<h3><a href="https://colab.research.google.com/drive/1cDXzCnTaVDEe3Dnhafp-Rdu6rKs3BuXI" target="_blank" rel="noopener nofollow ugc">Google Colaboratory</a></h3>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p><strong>Link to Colab Notebook</strong>. Only predicts kidneys.</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://colab.research.google.com/drive/1_49DjRyQhnlgB5z_lIyXkV5pYydEYvd3">
  <header class="source">
      <img src="https://ssl.gstatic.com/colaboratory-static/common/f6c736cbde16632a9fcd16d9ab1970b1/img/favicon.ico" class="site-icon" width="128" height="128">

      <a href="https://colab.research.google.com/drive/1_49DjRyQhnlgB5z_lIyXkV5pYydEYvd3" target="_blank" rel="noopener nofollow ugc">colab.research.google.com</a>
  </header>

  <article class="onebox-body">
    <img width="260" height="260" src="https://colab.research.google.com/img/colab_favicon_256px.png" class="thumbnail onebox-avatar">

<h3><a href="https://colab.research.google.com/drive/1_49DjRyQhnlgB5z_lIyXkV5pYydEYvd3" target="_blank" rel="noopener nofollow ugc">Google Colaboratory</a></h3>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


tom: <p>Hi Akshay,</p>
<p>I’d probably look at two diagnostic (ha) questions:</p>
<ul>
<li>If you permute the labels (<code>labels = labels.flip(1)</code> or so), does it still happen this way or can you then predict another organ?</li>
<li>If you stick with a single batch, can you overfit?</li>
</ul>
<p>The former is a check whether there is something wrong in the tensor shapes or somesuch. I tend to print all shapes to debug that type of thing as a first step.</p>
<p>The second would go in the direction of capacity.</p>
<p>Best regards</p>
<p>Thomas</p>

aksg87: <p>Thanks for the feedback <a class="mention" href="/u/tom">@tom</a>!</p>
<p>I reduced the resolution of the input and that fixed the issue suggesting that it was a capacity problem. I had checked permuting the labels and that still had the same behavior.</p>