Calculate the loss and dice for multi-class segmentation

sidra_Aleem: <p>I am working with multi-class segmentation. I have 4 classes, my input to model has dimesnion : <code>32,1,384,384</code>. The ground truth dimension is <code>32,4,384,384</code>. The prediction from the model has the dimension <code>32,4,384,384</code>. I am trying to calculate the loss using cross-entropy loss as :<code> loss = CE_loss(preds, torch.argmax(var_gt, dim=1))</code> (I want to use this specific loss as I am replicating a paper and authors used it). However, the dice score remains constant for both train and val set after 2nd epoch as shown below:  <div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/6/6682149f0d5f7fc77cbd8c362911325caaff3e67.png" data-download-href="https://discuss.pytorch.org/uploads/default/6682149f0d5f7fc77cbd8c362911325caaff3e67" title="learning curves"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/6/6682149f0d5f7fc77cbd8c362911325caaff3e67_2_690x182.png" alt="learning curves" data-base62-sha1="eCPoGZZASazvgdQfuUFRGa3rufR" width="690" height="182" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/6/6682149f0d5f7fc77cbd8c362911325caaff3e67_2_690x182.png, https://discuss.pytorch.org/uploads/default/original/3X/6/6/6682149f0d5f7fc77cbd8c362911325caaff3e67.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/6/6/6682149f0d5f7fc77cbd8c362911325caaff3e67.png 2x" data-dominant-color="FCFCFC"><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">learning curves</span><span class="informations">719×190 13.1 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>To perform diagnostic, I followed Andrej Karpathy’s tip of overfitting the model with train data using 2 images. I did that, but my loss was not 0, it came out to be 1.6575. I am not able to figure out what is the issue here. Can anayone please guide me. Below is also the code for dice calcualtion, but I beleive the main problem here is loss.Below is given the code for how I am calculating the dice score:</p>
<pre><code>eps = 0.0001
iflat = input_.reshape(-1)
tflat = target.reshape(-1)

if n_outputs == 1:
    intersection = (iflat * tflat).sum()
    union = iflat.sum() + tflat.sum()
    dice = (2.0 * intersection + eps) / (union + eps)
    if print_:
        print(dice)
else:
    dice = np.zeros(n_outputs-1)
    for c in range(1,n_outputs):   # assumes background is first class and doesn't compute its score
        iflat_ = iflat==c
        tflat_ = tflat==c
        intersection = (iflat_ * tflat_).sum()
        union = iflat_.sum() + tflat_.sum()
        d = (2.0 * intersection + eps) / (union + eps)
        if print_:
            print(c, d)
        dice[c-1] += d
    if average:
        dice = np.sum(dice)/(n_outputs-1)
</code></pre>

sidra_Aleem: <p>Since the new users can only put one embeded image i n the post, I am adding the results following Andrej’s tip here:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/7/d7b6975c6a6c5c0cab076a08709c2beccb24d1a1.png" data-download-href="https://discuss.pytorch.org/uploads/default/d7b6975c6a6c5c0cab076a08709c2beccb24d1a1" title="Results:"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/7/d7b6975c6a6c5c0cab076a08709c2beccb24d1a1_2_690x155.png" alt="Results:" data-base62-sha1="uMhR4RvW6DSCz7BBdFZG9CCiYTf" width="690" height="155" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/7/d7b6975c6a6c5c0cab076a08709c2beccb24d1a1_2_690x155.png, https://discuss.pytorch.org/uploads/default/optimized/3X/d/7/d7b6975c6a6c5c0cab076a08709c2beccb24d1a1_2_1035x232.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/d/7/d7b6975c6a6c5c0cab076a08709c2beccb24d1a1.png 2x" data-dominant-color="895C86"><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">Results:</span><span class="informations">1062×239 43.4 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>

ptrblck: <p>In your example code you are using some numpy operations, which could break the computation graph. Could you replace these using PyTorch operations and did you check the gradients in the model?</p>

sidra_Aleem: <p>How dice calcualtion could break the computation graph? The graidents are updated on the basis of loss, while dice score is the evaluation critertion to save the best model checkpoint.</p>
<pre><code class="lang-auto">val_dice =  dice_score(val_segmented_volume, gt_samples, n_outputs=n_channels_out)
if avg_val_dice &gt; best_acc:
                best_acc = avg_val_dice
                print("best_acc- after updation", best_acc)
                save_model(model, config, suffix, folder_time)
</code></pre>
<p>How to check the gradients in the model, I am not very expert in this field.</p>

ptrblck: <p>Using numpy would detach the tensor from the computation graph as seen here:</p>
<pre><code class="lang-python">input_ = torch.randn(10, 10, requires_grad=True)
target = torch.randint(0, 2, (10, 10)).float()
n_outputs = input_.size(0)

eps = 0.0001
iflat = input_.reshape(-1)
tflat = target.reshape(-1)

dice = np.zeros(n_outputs-1)
for c in range(1,n_outputs):   # assumes background is first class and doesn't compute its score
    iflat_ = iflat==c
    tflat_ = tflat==c
    intersection = (iflat_ * tflat_).sum()
    union = iflat_.sum() + tflat_.sum()
    d = (2.0 * intersection + eps) / (union + eps)
    print(c, d)
    dice[c-1] += d

print(dice)
# [1.92307311e-06 1.00000000e+00 1.00000000e+00 1.00000000e+00
 # 1.00000000e+00 1.00000000e+00 1.00000000e+00 1.00000000e+00
 # 1.00000000e+00]
</code></pre>
<p>Even though <code>input_</code> required gradients, the <code>dice</code> object is now a numpy array without any gradient history.<br>
I guess you are recreating a tensor afterwards and set its <code>.requires_grad</code> attribute to <code>True</code> again as a workaround for this issue, which won’t reattach the tensor to the computation graph.</p>

sidra_Aleem: <p>I would update numpy operation with torch. But for gradients I am still not clear what is your point. I am <strong>not</strong> updating the model based on dice score stats.  I previously added only chunk for dice calculation, but before that in <code>model.train()</code> mode, I do calculate the loss and backpropagate the error based on that as follows:</p>
<pre><code class="lang-auto">for i, batch in enumerate(train_loader):
        
            input_samples, gt_samples, voxel_dim = batch

            if torch.cuda.is_available():
                var_input = input_samples.to(device)
                var_gt = gt_samples.to(device)
          
            preds = model(var_input)
            if n_channels_out == 1:
               loss = weighted_cross_entropy_with_logits(preds, var_gt)
            else:
             
                loss = CE_loss(preds, torch.argmax(var_gt, dim=1))          

            train_loss_total += loss.item()
            optimizer.zero_grad()
            loss.backward() 
            optimizer.step()

</code></pre>
<p>Which means that in this step the gradients have been updated in the model. After that I caculate train dice. And to evaluate the model, I used validation set in <code>model.eval()</code> mode and calculated predictions, loss and dice score. And then saved the best checkpoint on the basis of dice score.</p>

ptrblck: <p>Thanks for clarifying. In that case I have misunderstood the question and don’t understand why you would expect to see a perfect dice score while overfitting two samples for a crossentropy loss.</p>

sidra_Aleem: <p>Apologies for confuion! I was trying to debug the issue that is why I was overfitting the model with two samples to see weather the loss eventrually goes to 0. My main issue is that, the dice score remain constant to 0.667 after the second epoch for both train and val set as shown in<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/a/9aae2d7fb0bf3afd2969c68dd04746809acd5e5a.png" data-download-href="https://discuss.pytorch.org/uploads/default/9aae2d7fb0bf3afd2969c68dd04746809acd5e5a" title="Leaning curve"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/a/9aae2d7fb0bf3afd2969c68dd04746809acd5e5a_2_634x500.png" alt="Leaning curve" data-base62-sha1="m4mHG9nGcMhhw8I005g6NQ9WSoO" width="634" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/a/9aae2d7fb0bf3afd2969c68dd04746809acd5e5a_2_634x500.png, https://discuss.pytorch.org/uploads/default/original/3X/9/a/9aae2d7fb0bf3afd2969c68dd04746809acd5e5a.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/9/a/9aae2d7fb0bf3afd2969c68dd04746809acd5e5a.png 2x" data-dominant-color="FCFDFD"><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">Leaning curve</span><span class="informations">695×548 28.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>.<br>
As per the learning curves model is learning, as loss is changing but it is not changing the dice accordingly. Should I be using softmax to un-normalized logits? Though the official documentation says that:<code> The input is expected to contain the unnormalized logits for each class (which do not need to be positive or sum to 1, in general).</code></p>

ptrblck: <p>You could visualize the predictions and try to estimate the dice as a quick smoke test.<br>
The issue I’m seeing is that you expect to see a reduced dice loss without optimizing for it and assume your <code>nn.CrossEntropyLoss</code> is a good proxy for it.<br>
Using a tiny subset of the data makes sense to make sure your model is able to overfit your currently used loss function, but I would not expect to see the same effect on other losses.<br>
The same applies for e.g. <code>nn.MSELoss</code>. I would also not expect to see a perfect overfitting of it in your setup.</p>