Indices from `max` is not differentiable in custom loss function

UserLee: <p>I’m training a network using a custom loss function. And it gives a <code>RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn</code>. It seems that the indices from <code>max</code> is not differentiable which i have to use for loss function.<br>
Does anyone have idea to fix it?<br>
This is my code:</p>
<pre><code class="lang-auto">class MyLoss(nn. Module):
    def __init__(self):
        super(MyLoss, self).__init__()
        
    def forwad(self, d, m_label, n_label):
        # d size : [N, C, H, W], from softmax output
        # m_label size : [N, 1, H, W]
        # n_label size : [N, 1, H, W]
        
        # get the predicted class
        pred = d.max(1)[1]
        pred = pred.unsqueeze(1)
        n_pred = 2.7*pred + m_label
        
        # calc image grad along x, y
        m_dx = m_label[:, :, 1:, :] - m_label[:, :, :-1, :]
        m_dy = m_label[:, :, :, 1:] - m_label[:, :, :, :-1]
        n_dx = n_pred[:, :, 1:, :] - n_pred[:, :, :-1, :]
        n_dy = n_pred[:, :, :, 1:] - n_pred[:, :, :, :-1]
        
        # loss1: minimize the difference of grad of m_label and n_pred
        loss1 = torch.mean(((n_dx - m_dx).abs().sum() 
                              + (n_dy - m_dy).abs().sum()))
        
        # loss2: minimize the difference of n_pred and n_label
        loss2 = nn.functional.l1_loss(n_pred, n_label)
        return loss1 + loss2
</code></pre>

KFrank: <p>Hi Lee!</p>
<aside class="quote no-group" data-username="UserLee" data-post="1" data-topic="88194" data-full="true">
<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/u/e0b2c6/48.png" class="avatar"> UserLee:</div>
<blockquote>
<p>It seems that the indices from <code>max</code> is not differentiable which i have to use for loss function.</p>
</blockquote>
</aside>
<p>First the “why” of your error:</p>
<p>The pytorch framework uses various gradient-descent methods<br>
to train your network – that is to minimize your loss function<br>
with respect to your model’s parameters.  For this to work, your<br>
loss function must be a (usefully) differentiable function of the<br>
model parameters.</p>
<p>As you have seen, the <code>argmax()</code> function (“the indices from<br>
max”) is not differentiable, so you can’t use it in your loss function.</p>
<p>As to how to fix your issue:</p>
<p>The short answer is that you have to modify your loss function<br>
to be (usefully) differentiable.</p>
<p>I don’t understand the concept behind your loss functions, so I<br>
don’t have suggestions for you.</p>
<p>What type of problem are you trying to solve?</p>
<p>What are the inputs to your network, and what do those inputs<br>
mean?</p>
<p>What are the meanings of <code>m_label</code> and <code>n_label</code>?</p>
<p>What, conceptually, determines whether your model has made<br>
a good prediction, and how is your loss function trying to measure<br>
that?</p>
<p>Good luck!</p>
<p>K. Frank</p>

UserLee: <p>Thanks for reply.<br>
I’m trying to solve a muti classification problem called phase unwrapping using a u-net. The final layer is 1x1 conv followed by softmax.<br>
There is a relationship between image <code>m_label</code> and <code>n_label</code>, <code>n_label = m_label + constant*k</code>, in which in which m_label is network input ,  k is predicted class (an integer matrix) from softmax output.<br>
To make a good prediction, i minimize:</p>
<ol>
<li>the diffference of grad of m_label and n_pred</li>
<li>the difference of n_label and n_pred (calc by <code>n_pred = m_label + constant*k_pred</code>)</li>
<li>the cross entry loss between k_pred and k_label (calc by <code>round((n_label - n_label)/constant)</code> ), not mentioned in code</li>
</ol>