UNet Model not learning for multiclass image segmentation task

ankita.buntolia: <p><strong>Problem: UNet model is not learning even after100 epochs. It shows the same loss for both training and validation.</strong></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/c/dc6f1fcba71e25179d4292913ef81e40584a3ff5.png" data-download-href="https://discuss.pytorch.org/uploads/default/dc6f1fcba71e25179d4292913ef81e40584a3ff5" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/d/c/dc6f1fcba71e25179d4292913ef81e40584a3ff5.png" alt="image" data-base62-sha1="vs37OlMVgEbMCLg1lIgZrO40mup" width="396" height="500" data-dominant-color="484848"><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">636×802 9.38 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><strong>Background information:</strong><br>
Each pixel of the Image can have only one label  (in total 4 classes as 0,1,2, or 3).<br>
From the DataSet:<br>
BEFORE TRANSFORMATION:<br>
→ label shape from OCTDataset: (496, 512)<br>
→ img shape from OCTDataset: (496, 512)<br>
→ image dtype from OCTDataset: float32, label dtype from OCTDataset: float32</p>
<p>AFTER TRANSFORMATION:<br>
→ shape just after transform of the label: torch.Size([1, 512, 512])<br>
–&gt;shape just after transform of the img: torch.Size([1, 512, 512])<br>
–&gt;image dtype after transform: torch.float32, label dtype after transform: torch.float32</p>
<p><strong>My UNet Model:</strong></p>
<pre><code class="lang-auto">#####  U-Net Model
## 
class DoubleConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(DoubleConv, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 3, 1, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, 3, 1, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
        )

    def forward(self, x):
        return self.conv(x)

class UNET(nn.Module):
    def __init__(
            self, in_channels=1, out_channels=3, features=[64, 128, 256, 512],
    ):
        super(UNET, self).__init__()
        self.ups = nn.ModuleList()
        self.downs = nn.ModuleList()
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

        # Down part of UNET
        for feature in features:
            self.downs.append(DoubleConv(in_channels, feature))
            in_channels = feature

        # Up part of UNET
        for feature in reversed(features):
            self.ups.append(
                nn.ConvTranspose2d(
                    feature*2, feature, kernel_size=2, stride=2,
                )
            )
            self.ups.append(DoubleConv(feature*2, feature))

        self.bottleneck = DoubleConv(features[-1], features[-1]*2)
        self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)
        # softmax
        self.softmax=nn.Softmax(dim=1)

    def forward(self, x):
      skip_connections = []

      for down in self.downs:
        x = down(x)
        skip_connections.append(x)
        x = self.pool(x)

      x = self.bottleneck(x)
      skip_connections = skip_connections[::-1]

      for idx in range(0, len(self.ups), 2):
        x = self.ups[idx](x)
        skip_connection = skip_connections[idx//2]

        if x.shape != skip_connection.shape:
          x = TF.resize(x, size=skip_connection.shape[2:])

        concat_skip = torch.cat((skip_connection, x), dim=1)
        x = self.ups[idx+1](concat_skip)

      return self.softmax(self.final_conv(x))
</code></pre>
<p><strong>My code for training:</strong></p>
<pre><code class="lang-auto"># Training function
def train(model, loader, optimizer, loss_fn, device):
    epoch_loss = 0.0

    model.train()
    for x, y in loader:
      x = x.to(device, dtype=torch.float32)
      y = y.to(device, dtype=torch.float32)        

      optimizer.zero_grad()
      y_pred = model(x)
      # removing channel from orig label
      y = y.squeeze(1)

      loss = loss_fn(y_pred, y.type(torch.LongTensor).cuda())
      loss.backward()
      optimizer.step()
      epoch_loss += loss.item()* x.size(0)
    epoch_loss = epoch_loss/len(loader)
    return epoch_loss

# Evaluation function
def evaluate(model, loader, loss_fn, device):
    epoch_loss = 0.0

    model.eval()
    with torch.no_grad():
      for x, y in loader:
        x = x.to(device, dtype=torch.float32)
        y = y.to(device, dtype=torch.float32)

        y_pred = model(x)

        # removing channel from orig label
        y = y.squeeze(1)

        loss = loss_fn(y_pred, y.type(torch.LongTensor).cuda())
        epoch_loss += loss.item()* x.size(0)     

      epoch_loss = epoch_loss/len(loader)

    return epoch_loss
</code></pre>
<pre><code class="lang-auto">""" Hyperparameters """
H = 512
W = 512
size = (H, W)
num_epochs = 100
lr = 1e-4 # 0.001
device = torch.device('cuda')   ## 

""" Calculate the time taken """
def epoch_time(start_time, end_time):
    elapsed_time = end_time - start_time
    elapsed_mins = int(elapsed_time / 60)
    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))
    return elapsed_mins, elapsed_secs

class_weight = [0.9851726793996959,  0.005821095764414637,  0.005777171078234721,  0.0032290537576547757] # class weights of whole dataset (before making subset of training, validation, and test dataset)
class_weight_tensor = torch.Tensor(class_weight).to(device, dtype=torch.float32)


model = UNET(in_channels=1, out_channels=4)
model = model.to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=lr)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=5, verbose=True)
# loss function 
loss_fn = nn.CrossEntropyLoss(weight=class_weight_tensor)


# lists to collects data while treaining
best_valid_loss = float("inf")
train_losses_list = list()

# if you want to resume training from check point then change it to True
resume_Training = False
</code></pre>
<p><strong>Training loop</strong></p>
<pre><code class="lang-auto">""" Training the model """
for epoch in range(num_epochs):
  # resume
  if resume_Training:
    checkpoint = torch.load('/content/drive/MyDrive/Practical_work/May2022/outputs/model_V2.pth')
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    epoch = checkpoint['epoch']
    train_losses_list = checkpoint['loss_train_list']
    val_losses_list = checkpoint['loss_val_list']
    best_valid_loss = val_losses_list[-1]
    model.train()
    resume_Training = False


  start_time = time.time()


  train_loss = train(model, oct_trainingloader, optimizer, loss_fn, device = device)
  valid_loss = evaluate(model, oct_validationloader, loss_fn, device = device) # epoch_loss, precision, recall, thresholds, iou_val

  # append 
  train_losses_list.append(train_loss)
  val_losses_list.append(valid_loss)



  """ Saving the model """
  if valid_loss &lt; best_valid_loss:
    data_str = f"Valid loss improved from {best_valid_loss:2.4f} to {valid_loss:2.4f}. Saving checkpoint"
    print(data_str)

    best_valid_loss = valid_loss
    state = {'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss_train_list': train_losses_list,
            'loss_val_list': val_losses_list,
            }
    torch.save(state, f'/content/drive/MyDrive/Practical_work/May2022/outputs/model_V2.pth')        

  end_time = time.time()
  epoch_mins, epoch_secs = epoch_time(start_time, end_time)

  data_str = f'Epoch: {epoch+1:02} | Epoch Time: {epoch_mins}m {epoch_secs}s'
  data_str += f'\tTrain Loss: {train_loss:.3f}'
  data_str += f'\t Val. Loss: {valid_loss:.3f}'

  print(data_str)


# Ploting losses
plt.plot(train_losses_list)
plt.plot(val_losses_list)```</code></pre>

ptrblck: <p><code>nn.CrossEntropyLoss</code> expects raw logits as the model outputs, not probabilities, so remove the <code>nn.Softmax</code> from your model and rerun the training.</p>

ankita.buntolia: <p>I tried after removing softmax from UNet model. But the training does not seem right. I tried with 30 epochs just to cross-check, but in the second epoch, loss dropped from 36% to 5%, which does not seem right.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/4/a/4a9b3e5c5de728dd695c269504c1c31199556f5b.png" alt="image" data-base62-sha1="aDZWjZIVvRVQjWBvxcCXz9IMOtJ" width="379" height="248"></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/e/6ef920ff61d24bc8b50f432e166f8b214fa5b9e0.png" data-download-href="https://discuss.pytorch.org/uploads/default/6ef920ff61d24bc8b50f432e166f8b214fa5b9e0" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/6/e/6ef920ff61d24bc8b50f432e166f8b214fa5b9e0.png" alt="image" data-base62-sha1="fPIher4sSeChLiNma055ORjXKH6" width="354" height="500" data-dominant-color="ECEBEC"><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">604×853 13.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>
<p>Labels are highly imbalanced, so I used weights in CE loss. I calculated my weights as shown here:<br>
Classes → 0 (background), and remaining three (1,2, and 3) are class indices in labels.</p>
<p>First I consider whole labels from the whole dataset (dataset before splitting into training, validation, and testing) and counted how many times each class appeared. With these counts, I calculated weights.</p>
<pre><code class="lang-auto">class_counts = [294219409,   1738456,   1725338,    964349]
total_count = sum(class_counts)
class_weight = [class_counts[0]/total_count, class_counts[1]/total_count, class_counts[2]/total_count, class_counts[3]/total_count ]
class_weight
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/f/ffe57abe98d27b5ef2a8bfef928b411c36f96212.png" data-download-href="https://discuss.pytorch.org/uploads/default/ffe57abe98d27b5ef2a8bfef928b411c36f96212" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/f/f/ffe57abe98d27b5ef2a8bfef928b411c36f96212.png" alt="image" data-base62-sha1="AvLpzvw7C1vihdOObJD94O8WCeC" width="690" height="110" data-dominant-color="F3F4F4"><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">1110×178 8.01 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>Later I fed this weights to CE Loss:</p>
<pre><code class="lang-auto">loss_fn = nn.CrossEntropyLoss(weight=class_weight_tensor)
</code></pre>
<p>Is this the correct way to calculate class weights?</p>
<p>Currently, I am only using CE Loss with class weights for training. I know Dice loss and Focal Loss would be ideal one but even with just CE Loss network should learn and work (even with not so good efficiency as compared with Dice Loss).</p>
<p>Do you see any mistakes in my process?</p>
<p>Thank you very much for your support!</p>
<p>Best Regards,<br>
Ankita</p>

ptrblck: <aside class="quote no-group" data-username="ankita.buntolia" data-post="3" data-topic="160656">
<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/bc79bd/48.png" class="avatar"> ankita.buntolia:</div>
<blockquote>
<p>I tried with 30 epochs just to cross-check, but in the second epoch, loss dropped from 36% to 5%, which does not seem right.</p>
</blockquote>
</aside>
<p>The loss is not given in a percentage and an initial drop is common, so could you explain your concern a bit?</p>
<aside class="quote no-group" data-username="ankita.buntolia" data-post="3" data-topic="160656">
<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/bc79bd/48.png" class="avatar"> ankita.buntolia:</div>
<blockquote>
<p>Is this the correct way to calculate class weights?</p>
</blockquote>
</aside>
<p>No, since you are now weighting the majority class higher and are thus adding more penalty to a wrong classification of class0. Divide by the class frequencies as explained in <a href="https://discuss.pytorch.org/t/predictions-stuck-at-zero-when-positive-label-1-is-only-16-of-data/145203/4">this post</a>.</p>

ankita.buntolia: <aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="160656">
<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>The loss is not given in a percentage and an initial drop is common, so could you explain your concern a bit?</p>
</blockquote>
</aside>
<p>I mean with 0.36 to 0.05!</p>
<aside class="quote no-group quote-modified" data-username="ptrblck" data-post="4" data-topic="160656">
<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>No, since you are now weighting the majority class higher and are thus adding more penalty to wrong classification of class0. Divide by the class frequencies as explained in <a href="https://discuss.pytorch.org/t/predictions-stuck-at-zero-when-positive-label-1-is-only-16-of-data/145203/4">this post </a>.</p>
</blockquote>
</aside>
<p>Thanks for the link. It’s very helpful. Based on this link, I have calculated weights per batch for Cross entropy loss. Still, with training 100 epochs, the learning is quite slow.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/5/7/57224faa33fe5fe7112f0d2415c87acadf35a140.png" alt="image" data-base62-sha1="cqP3bg8DBMJkLOWF5PBE5E8YJvG" width="378" height="248"></p>
<p>I saw many code repo where they use a combination of Cross entropy loss and Dice Score, but the contribution of each one of them in final total loss varies (<a href="https://github.com/ai-med/relaynet_pytorch/blob/40ae1aa56e426da14ddca37e06c2f31966febea5/build/lib/relaynet_pytorch/net_api/losses.py" rel="noopener nofollow ugc">some use just addition of both the loss</a>, and  <a href="https://github.com/usuyama/pytorch-unet/blob/master/pytorch_unet.ipynb" rel="noopener nofollow ugc">some use percentages of (50%) from each one of them</a>). Are there any useful guidelines while calculating the loss function especially for imbalanced dataset?</p>

Coca_Cola12: <aside class="quote no-group" data-username="ankita.buntolia" data-post="3" data-topic="160656">
<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/bc79bd/48.png" class="avatar"> ankita.buntolia:</div>
<blockquote>
<p><code>38,    </code></p>
</blockquote>
</aside>
<p>Hello! im having the same issue. Did you solve this problem ankita?</p>