My custom loss never goes backward

username: tongliuTL
cooked: <p>I am implementing a custom loss for my multi-class classification application.<br>
For every y_true, there is a corresponding P_true which represents a list of values.<br>
For every y_pred, after some reformat, using it as the key of an external dictionary, can look up another list of values (with equal length of P_true).<br>
I’d like to calculate the MSE loss between these two lists of values (true_samples, pred_samples), and use the result as the output of this custom loss function.</p>
<pre><code class="lang-auto">def custom_loss_fn(y_pred, y_true, P_true):
    y_pred_softmax = torch.log_softmax(y_pred, dim=1)
    _, y_pred_tags = torch.max(y_pred_softmax, dim=1)

    true_samples = Variable(P_true, requires_grad=True).to(device)

    pred_cluster_names = [idx2class[tag.squeeze().tolist()]
                          for tag in y_pred_tags]
    pred_samples = Variable(
        torch.FloatTensor(
            [name_samples_mapping[pred_name]
                for pred_name in pred_cluster_names]
        ), requires_grad=True).to(device)

    output = mseloss(true_samples, pred_samples)

    return output
</code></pre>
<p>But I found that the training loss from each epoch never changed, suggesting something wrong with my computation graph connectivity that autograd failed. Could anyone provide help? Thanks!</p>
<p>Edit: true_samples tensors have <code>requires_grad=True</code>, pred_samples tensors have <code>grad_fn=&lt;CopyBackwards&gt;</code>, and output for training loss has <code>grad_fn=&lt;MeanBackward0&gt;</code>.</p>

username: albanD
cooked: <p>Hi,</p>
<p>The problem is that the argmax operation that you do on <code>y_pred</code> at the beginning is not differentiable. It returns integer indices.</p>
<p>So no gradient will be able to flow back to <code>y_pred</code> and to your network above <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=9" title=":confused:" class="emoji" alt=":confused:"></p>

username: tongliuTL
cooked: <p>Thank you for your insight! I am new to PyTorch and didn’t realize <em>argmax</em> is the root cause of the problem and have been doubting if my dictionary lookups broke the connectivity since they were not wrapped in <code>Variable</code>s.</p>
<p>I read your answer to another similar problem (<a href="https://discuss.pytorch.org/t/build-your-own-loss-function-in-pytorch/235/65">link here</a>), would you mind providing some further suggestions/solutions for work around this argmax issue and make my backprop working? Thanks!</p>

username: albanD
cooked: <p>As mentioned in the other thread, replacing this kind of argmax is a area of research in itself almost <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>I guess the simplest you can try here is to use a weighted version based on the softmax value. So that the predicted value is a weighted value of different dictionary entries.<br>
But this might not work…</p>

username: tongliuTL
cooked: <p>I think essentially my custom loss function is more like a <em>ranking</em>-purpose function than a <em>classification</em> one. Saying this I mean: if the <code>y_pred</code> looks up <code>pred_samples</code> which can satisfy less <code>mseloss</code> when being compared to <code>true_samples</code>, then <code>y_pred</code> is assumed to stand out in the candidate list among all possible <code>y</code> classes.</p>
<p>Also, I read some other resources, do you think if these materials can lead me in the right direction?</p>
<ul>
<li><a href="https://bouthilx.wordpress.com/2013/04/21/a-soft-argmax/" rel="nofollow noopener">https://bouthilx.wordpress.com/2013/04/21/a-soft-argmax/</a></li>
<li><a href="https://github.com/MWPainter/cvpr2019/blob/master/stitched/soft_argmax.py#L82" rel="nofollow noopener">https://github.com/MWPainter/cvpr2019/blob/master/stitched/soft_argmax.py#L82</a></li>
</ul>
<p>Thanks a lot for your time and reply.</p>