My model is predicting everything as background

MA_CASANDRA_QUILANG: <p>Hi! I am currently working to develop a model that can hopefully take an RGB image containing garbage and outputs a segmented image (i am hoping to segment it according to the 60 categories I defined). I used Adam as the optimizer and cross-entropy as the loss_function. During training, (I tried 10 epoch at most), the lowest loss value I got was 0.14. This is my training script:</p>
<pre><code class="lang-auto">BATCH_SIZE = 4
EPOCHS = 10

def train(model):
  model.train()
  for epoch in range(EPOCHS):
      for i in tqdm(range(0, len(img_train), BATCH_SIZE)): 
          batch_img_train = img_train[i:i+BATCH_SIZE].view(-1, 3, 448, 448)
          batch_mask_train = mask_train[i:i+BATCH_SIZE].view(-1, 1, 448, 448)
        
          model.zero_grad()

          outputs = model(batch_img_train)
        
          loss = loss_function(outputs, batch_mask_train.squeeze(1).long())
          loss.backward()
          optimizer.step()    # Does the update

      print(f"Epoch: {epoch}. Loss: {loss}")
  
  return batch_img_train, batch_mask_train, outputs

train(model)

</code></pre>
<p>while this is my validation script ( I also tried to print predictions here):</p>
<pre><code class="lang-auto">def test(model):
  model.eval
  correct = 0
  total = 0

  with torch.no_grad():
      for i in tqdm(range(len(img_test))):

          
          real_class = mask_test[20]
          net_out = model(img_test[i].view(-1, 3, 224, 224))[0]
          predicted_class = torch.argmax(net_out, 0)
          
          prediction = predicted_class.eq(mask_test[20])

          
          plt.figure(figsize=(22,8))

          
          plt.subplot(1, 3, 1)
          plt.imshow(img_test[20].view(-1, 224, 224, 3).detach().cpu().squeeze())
          plt.title('original image')

          plt.subplot(1, 3, 2)
          plt.imshow(real_class.cpu().squeeze())
          plt.title('ground truth')

          plt.subplot(1, 3, 3)
          plt.imshow(predicted_class.cpu().squeeze())
          plt.title('predicted mask')
          

      
test(model)
</code></pre>
<p>but my predictions are way to false because my model tends to classify every pixel as backgroudn despite having a small loss value. I read about class imbalance and even tried implementing IoU in the training but I couldnt get it right. I am very confused right now, can someone please please help me?</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/e/7e9f2e61b6acbd59e5bebcc1501088abe6ec9fb8.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/7e9f2e61b6acbd59e5bebcc1501088abe6ec9fb8" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/e/7e9f2e61b6acbd59e5bebcc1501088abe6ec9fb8_2_690x221.jpeg" alt="image" data-base62-sha1="i49cJtttzoCmJ686vZpVkBTIsGI" width="690" height="221" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/e/7e9f2e61b6acbd59e5bebcc1501088abe6ec9fb8_2_690x221.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/7/e/7e9f2e61b6acbd59e5bebcc1501088abe6ec9fb8_2_1035x331.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/7/e/7e9f2e61b6acbd59e5bebcc1501088abe6ec9fb8_2_1380x442.jpeg 2x" data-dominant-color="7B557B"><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">1444×463 111 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 Cassie!</p>
<aside class="quote no-group" data-username="MA_CASANDRA_QUILANG" data-post="1" data-topic="108701" data-full="true">
<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/ma_casandra_quilang/48/13913_2.png" class="avatar"> MA_CASANDRA_QUILANG:</div>
<blockquote>
<p>my predictions are way to false because my model tends to classify every pixel as backgroudn despite having a small loss value. I read about class imbalance</p>
</blockquote>
</aside>
<p>Class imbalance is likely an issue you will need to deal with.  Have<br>
you tried the <code>weight</code> argument passed to the <code>CrossEntropyLoss</code><br>
constructor?  One typically weights each class in inverse proportion<br>
to its frequency of appearance.  (If you are using the functional version<br>
of cross entropy, <code>torch.nn.functional.cross_entropy()</code>, it also<br>
has such a <code>weight</code> argument.)</p>
<aside class="quote no-group">
<blockquote>
<p>(i am hoping to segment it according to the 60 categories I defined).</p>
</blockquote>
</aside>
<p>One experiment you could try is to combine your 60 different categories<br>
into one “garbage” class, and train a two-class problem: background<br>
(i.e., not garbage) vs. garbage.  This will reduce the class imbalance,<br>
because your garbage class will occur much more frequently than<br>
any individual one of your 60 classes.  Of course, you may still have a<br>
background / garbage imbalance, and may still want to use the <code>weight</code><br>
argument.</p>
<p>If you try this experiment, I would advise that you still treat this as a<br>
multi-class problem (that happens to have only two classes), and<br>
still use <code>CrossEntropyLoss</code>, rather than reformulating it as a binary<br>
classification problem (that uses something like <code>BCEWithLogitsLoss</code>).<br>
The idea is to get your current code working on a problem without as<br>
much class imbalance, only changing your code as much as necessary.</p>
<p>Good luck.</p>
<p>K. Frank</p>

MA_CASANDRA_QUILANG: <p>Thank you so much for your response! I will read  about your suggestion on using weight argument in my loss function since that idea is very new to me. I’ll get back here after I try this.</p>

MA_CASANDRA_QUILANG: <p>Hello!</p>
<p>I have tried using weighted cross entropy and my training script looks like this:</p>
<pre><code class="lang-auto">weights = [0.2939, 0.0105, 0.0005, 0.0015, 0.0002, 0.0015, 0.0551, 0.0181, 0.0453, 0.0135, 0.0037, 0.0049, 0.0024, 0.0370, 0.0007, 0.0193, 0.0022, 0.0100, 0.0100, 0.0066, 0.0007, 0.0152, 0.0223, 0.0027, 0.0010, 0.0000, 0.0017, 0.0010, 0.0169, 0.0015, 0.0419, 0.0012, 0.0007, 0.0027, 0.0167, 0.0056, 0.0000, 0.0759, 0.0010, 0.0049, 0.0054, 0.0120, 0.0007, 0.0086, 0.0022, 0.0010, 0.0091, 0.0032, 0.0005, 0.0010, 0.0071, 0.0174, 0.0069, 0.0029, 0.0015, 0.0015, 0.0269, 0.0010, 0.0191, 0.0659, 0.0556]
class_weights = torch.FloatTensor(weights)

import torch.optim as optim
optimizer = optim.Adam(model.parameters(), lr = 3e-5)
loss_CE = nn.CrossEntropyLoss(weight=class_weights)

BATCH_SIZE = 4
EPOCHS = 2


def train(model):
  model.train()
  for epoch in range(EPOCHS):
      for i in tqdm(range(0, len(img_train), BATCH_SIZE)): 
          batch_img_train = img_train[i:i+BATCH_SIZE].view(-1, 3, 112, 112)
          batch_mask_train = mask_train[i:i+BATCH_SIZE].view(-1, 1, 112, 112)
        
          model.zero_grad()

          outputs = model(batch_img_train)
          loss_ce= loss_CE(outputs, batch_mask_train.squeeze(1).long())
          loss = loss_ce
        
          loss.backward()
          optimizer.step()    # Does the update

      print(f"Epoch: {epoch}. Loss: {loss}")

train(model)

</code></pre>
<p>Am I doing it right?</p>

AlperEkmekci: <p>Hello,</p>
<p>What is your prediction output, you only get output class as result of prediction ? If your answer is no, doing <code>color mapping</code>  on your classes will solve your problem.</p>

MA_CASANDRA_QUILANG: <p>Hi! This is how my code for trying to show sample prediction:</p>
<pre><code class="lang-auto">def test(model):
  model.eval
  correct = 0
  total = 0

  with torch.no_grad():
      for i in tqdm(range(len(img_test))):

          
          real_class = mask_test[i]
          net_out = model(img_test[i].view(-1, 3, 224, 224))[0]
          _, predicted_class = torch.max(net_out, 0)
          prediction = predicted_class.eq(mask_test[i])
        

          
          plt.figure(figsize=(22,8))

          
          plt.subplot(1, 3, 1)
          plt.imshow(img_test[i].view(-1, 224, 224, 3).detach().cpu().squeeze())
          plt.title('original image')

          plt.subplot(1, 3, 2)
          plt.imshow(real_class.cpu().squeeze())
          plt.title('ground truth')

          plt.subplot(1, 3, 3)
          plt.imshow(predicted_class.cpu().squeeze())
          plt.title('predicted mask')
          
          break

      
test(model)
</code></pre>
<p>I have already tried using weighted cross entropy but the result was the same my model predicts everything as a background</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/b/3b4a5f74e25a80bcdf2a9dd7afc798b660544542.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/3b4a5f74e25a80bcdf2a9dd7afc798b660544542" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/3/b/3b4a5f74e25a80bcdf2a9dd7afc798b660544542.jpeg" alt="image" data-base62-sha1="8svvRtpJoiUA7cuVKlxS9fkxrhw" width="690" height="219" data-dominant-color="77517C"><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">715×227 27.3 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>

MA_CASANDRA_QUILANG: <p>how do I do color mapping?</p>

AlperEkmekci: <p>I assume that your prediction is consist more results except background prediction. Here is a snippet to colorize your output.</p>
<pre><code class="lang-auto">label_colors = np.array([(0, 0, 0), 
               # 0=background
               # 1=aeroplane, 2=bicycle, 3=bird, 4=boat, 5=bottle
               (128, 0, 0), (0, 128, 0), (128, 128, 0), (0, 0, 128), (128, 0, 128),
               # 6=bus, 7=car, 8=cat, 9=chair, 10=cow
               (0, 128, 128), (128, 128, 128), (64, 0, 0), (192, 0, 0), (64, 128, 0),
               # 11=dining table, 12=dog, 13=horse, 14=motorbike, 15=person
               (192, 128, 0), (64, 0, 128), (192, 0, 128), (64, 128, 128), (192, 128, 128),
               # 16=potted plant, 17=sheep, 18=sofa, 19=train, 20=tv/monitor
               (0, 64, 0), (128, 64, 0), (0, 192, 0), (128, 192, 0), (0, 64, 128)])

  r = np.zeros_like(image).astype(np.uint8)
  g = np.zeros_like(image).astype(np.uint8)
  b = np.zeros_like(image).astype(np.uint8)
  
  for l in range(0, nc):
    idx = image == l
    r[idx] = label_colors[l, 0]
    g[idx] = label_colors[l, 1]
    b[idx] = label_colors[l, 2]
    
  rgb = np.stack([r, g, b], axis=2)
  return rgb
</code></pre>
<p>This is a function to decode and plot your prediciton.</p>
<p>For more information about that you can check this <a href="https://learnopencv.com/pytorch-for-beginners-semantic-segmentation-using-torchvision/" rel="noopener nofollow ugc">link</a> out.</p>

MA_CASANDRA_QUILANG: <p>Thank you so much for your help. I will try this out!</p>

AlperEkmekci: <p>Waiting for you feedback, have a good day.</p>

MA_CASANDRA_QUILANG: <p>Hi! I tried the codes form the link you gave. I was stuck at this lines:</p>
<pre><code class="lang-auto">  real_class = mask_test[20]
          net_out = model(img_test[20].view(-1, 3, 224, 224))
          print(net_out.shape)
          predicted_class = torch.argmax(net_out.squeeze(), dim = 0).detach().cpu().numpy()
          print(predicted_class.shape)
          print(np.unique(predicted_class))
</code></pre>
<p>in the line where I tried printing np.unique, I should be getting the prediction already. Because in the link when they printed this certain code, the result was [0 3] and when I checked their color map that corresponds to a bird. But mine always prints out [0], which give push me suspecting that maybe my network indeed is just predicting everything as background. What should I do?</p>

MA_CASANDRA_QUILANG: <p>These are my results of printing:</p>
<pre><code class="lang-auto"> 0%|          | 0/150 [00:00&lt;?, ?it/s]torch.Size([1, 61, 224, 224])
(224, 224)
[0]</code></pre>

MA_CASANDRA_QUILANG: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/6/56cd10ff69a14dcb776bce657b1ccdab19a763c6.png" data-download-href="https://discuss.pytorch.org/uploads/default/56cd10ff69a14dcb776bce657b1ccdab19a763c6" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/5/6/56cd10ff69a14dcb776bce657b1ccdab19a763c6_2_690x461.png" alt="image" data-base62-sha1="cnSpJZ3gXsJXcv0TDqGJOHdBRI2" width="690" height="461" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/5/6/56cd10ff69a14dcb776bce657b1ccdab19a763c6_2_690x461.png, https://discuss.pytorch.org/uploads/default/original/3X/5/6/56cd10ff69a14dcb776bce657b1ccdab19a763c6.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/5/6/56cd10ff69a14dcb776bce657b1ccdab19a763c6.png 2x" data-dominant-color="C6C6C6"><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">797×533 16.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>

francois-rozet: <p>Hello <img src="https://discuss.pytorch.org/images/emoji/apple/wave.png?v=9" title=":wave:" class="emoji" alt=":wave:"> For <strong>binary</strong> segmentation with high class imbalance, I believe the Dice loss or Intersection over Union loss are better than the Cross Entropy.</p>
<p>Here is an implementation of both</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/francois-rozet/adopptrs/blob/4b680f72748033ffc37ab67926f5444b243e4922/python/criterions.py#L57-L85" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/francois-rozet/adopptrs/blob/4b680f72748033ffc37ab67926f5444b243e4922/python/criterions.py#L57-L85" target="_blank" rel="noopener nofollow ugc">francois-rozet/adopptrs/blob/4b680f72748033ffc37ab67926f5444b243e4922/python/criterions.py#L57-L85</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="57" style="counter-reset: li-counter 56 ;">
<li>class DiceLoss(nn.Module):</li><li>	'''Dice Loss (F-score, ...)'''</li><li></li><li>	def __init__(self, smooth=1.):</li><li>		super().__init__()</li><li></li><li>		self.smooth = smooth</li><li></li><li>	def forward(self, outputs, targets):</li><li>		inter = (outputs * targets).sum()</li><li>		dice = (2. * inter + self.smooth) / (outputs.sum() + targets.sum() + self.smooth)</li><li></li><li>		return 1. - dice</li><li></li><li></li><li>class IOULoss(nn.Module):</li><li>	'''Intersection Over Union Loss'''</li><li></li><li>	def __init__(self, smooth=1.):</li><li>		super().__init__()</li>
</ol></code></pre>

  This file has been truncated. <a href="https://github.com/francois-rozet/adopptrs/blob/4b680f72748033ffc37ab67926f5444b243e4922/python/criterions.py#L57-L85" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>Note that the mask should be 0 or 1 and the prediction a single number between 0 and 1 (not two like cross entropy).</p>

MA_CASANDRA_QUILANG: <p>hi thank you for taking time to reply. Can I still use dice loss if my class is more than 2? I have 60 classes plus background</p>

francois-rozet: <p>It is possible to compute IOU and Dice loss with more than two classes, but not with this implementation. Basically, you would compute the ratios for each class (<strong>except the background</strong>) individually before averaging them (weighted or not).</p>