Custon dice_loss function does not minimize the loss

username: joekid
cooked: <p>I am new to pytorch.  I’m working on semantic segmentation, so I like to use the dice_loss to update the model’s parameters (Previously, I tested the model with the CrossEntropy loss function and works reasonably).</p>
<p>Here, some code:</p>
<blockquote>
<p>for epoch in range(epochs):<br>
for inputs, labels in train_loader:<br>
inputs, labels = inputs.to(device), labels.to(device)<br>
optimizer.zero_grad()<br>
outputs = model(inputs)<br>
_, predicted = torch.max(outputs.data, 1) # since I will work with 1s and 0s<br>
loss = dice_loss(predicted,labels)<br>
loss.backward()  # Backward propagation<br>
optimizer.step()</p>
</blockquote>
<p>A loss output is as like tensor(0.7531, grad_fn=&lt;<strong>AddBackward0</strong>&gt;), and loss.backward() does not fails, but in my test the loss never decreases after it inits.</p>
<p>Here my loss function in details:</p>
<pre><code class="lang-auto">def dice_loss(predicted, labels):
    """Dice coeff loss for a batch"""
    
    # both the predicted and the labels data are being one-hot encoded
    onehot_pred = torch.Tensor()
    onehot_lab = torch.Tensor()    
    for batch, data in enumerate(zip(predicted, labels)):
        # to_categorical is the KERAS adapted function
        pred = utils.to_categorical(data[0]).unsqueeze(dim=0)
        lab = utils.to_categorical(data[1]).unsqueeze(dim=0)        
        onehot_pred = torch.cat((onehot_pred,pred),dim=0)
        onehot_lab = torch.cat((onehot_lab,lab),dim=0)
    
    # calculate the loss function
    ratio = 1 / predicted.size(0) # instead to divide by the batch_size
    
    # loss accumulator
    dc = torch.tensor(0).float()
    
    # I put required gradient in order to create a grad_fn
    # without loss.backward() does not work
    onehot_pred.requires_grad = True
    onehot_lab.requires_grad = True
    
    for batch , data in enumerate(zip(onehot_pred,onehot_lab)):
        dc += dice_coeff(data[0], data[1])*ratio
    return dc
</code></pre>
<p>and the other function</p>
<blockquote>
<p>def dice_coeff(predicted, labels):<br>
“”“Dice coeff for a simple plane”“”<br>
eps = 0.0001<br>
inter = torch.dot(predicted.view(-1), labels.view(-1))<br>
union = torch.sum(predicted) + torch.sum(labels) + eps<br>
dice_coeff = 1 - (2 * inter.float() + eps) / union.float()<br>
return dice_coeff</p>
</blockquote>

username: albanD
cooked: <p>Hi</p>
<p>The problem is that if you have to do add <code>onehot_pred.requires_grad = True</code>, that means it was False before. And so that means that some non-differentiable operation was applied to it.</p>
<p>Also fyi you can change <code>dc = torch.tensor(0).float()</code> to <code>dc = 0.</code><br>
And <code>onehot_pred = torch.Tensor()</code> to <code>onehot_pred = []</code>, <code>onehot_pred = torch.cat((onehot_pred,pred),dim=0)</code> to <code>onehot_pred.append(pred)</code> and add after the loop <code>onehot_pred = torch.cat(onehot_pred, dim=0)</code>.</p>

username: joekid
cooked: <p>I am seeing that the problem starts at</p>
<pre><code class="lang-auto">`_, predicted = torch.max(outputs.data, 1) `
</code></pre>
<p>I read that <strong>torch.max</strong> or <strong>torch.argmax</strong> are non differentiable functions. Until this line, <strong>outputs</strong>.grad_fn has backward and requires_grad is True, but <strong>predicted</strong> is  an object torch.return_types.max, so it doesn’t have grad_fn and it can no more propagate gradients backward.</p>

username: albanD
cooked: <p>Right.<br>
You should never use <code>.data</code>. It is both breaking the graph and unsafe.<br>
Max is a differentiable function where the gradient just flows back to the maximum value and all the other entries get 0.<br>
Argmax on the other hand is not differentiable as it return integer values. For which you can’t get gradients.</p>

username: joekid
cooked: <p>Hi <a class="mention" href="/u/alband">@albanD</a>, sorry for my delayed answer. I hope you help me. I modified my function in many ways (using Function, using Module,…), but the final computed <strong>loss</strong> is only a tensor, without any <strong>grad_fn</strong>  (That is, it is out the computation graph).</p>
<p>My goal is get dice_coeff and backward it (the formula isn’t exact).</p>
<p>Can you allow me to show line by line?</p>
<blockquote>
<p>criterion = DiceCoeffLoss() # custom loss function declaration, its code is below</p>
</blockquote>
<blockquote>
<p>outputs = model(inputs)</p>
</blockquote>
<p>and his gradient function:<br>
<code>outputs.grad_fn Out[285]: &lt;SigmoidBackward at 0x1b936013e08</code> # rigth</p>
<blockquote>
<p>loss = criterion(outputs,targets)</p>
</blockquote>
<p><strong>outputs</strong> shape is [n_batch, n_class, H, W]<br>
<strong>target</strong> shape is [n_batch,  H, W]</p>
<p>Into <code>criterion</code>, in order to calculate the loss, first I have to get the predicted labels along the channels:</p>
<blockquote>
<p>pred = torch.max(outputs, dim=1)</p>
</blockquote>
<p>Here, <code>pred</code> has no direct attribute grad_fn,  <code>pred.values </code> yet contains the  <code>outputs </code> values and <code>pred.values.grad_fn</code> has the gradient function <code> &lt;MaxBackward0 at 0x1b93601d1c8&gt;</code>.</p>
<p>However, since next operations is based on the predicted labels contained in  <code>pred.indices </code>, the final  <code>loss </code> is only a tensor.</p>
<p>Here the code (disclaimer, the formula isn’t the rigth):</p>
<pre><code class="lang-auto">def dice_coeff(predicted, targets):
    smooth = 1.0        
    inter = predicted.eq(targets).sum() # only true predited
    union = float(predicted.nelement() + targets.nelement() ) # all elements        
    dc = (2*inter + smooth) / (union + smooth)        
    return  dc

# I used nn.Module thinking it will add those operations to the graph    
class DiceCoeffLoss(nn.Module):
    def __init__(self):
        super(DiceCoeffLoss, self).__init__()        
         
    def forward(self, outputs, targets):
        
        pred = torch.max(outputs, dim=1)
        coeff = dice_coeff(pred.indices,targets)    
        return 1 - coeff
</code></pre>

username: albanD
cooked: <p>As I said above, the problem is that the argmax function is not differentiable. You will have to use something else.<br>
You can try to use soft attention where instead of the one hot encoding of the argmax, you use the result of the softmax layer.<br>
But that will depend on your exact problem what will work or not.</p>

username: joekid
cooked: <p>I resolved it. I learned the model output only should be transformed by some differentiable function (grad_fn) , such as exp, log, etc, in order to compute some loss value. Perhaps it was my lack of fundamentals of DL and Pytorch skills. Thank again.</p>