Torch.argmax() cause loss.backward() don't work

username: KirutoCode
cooked: <p>I am now training a network with Pytorch. In the training, I use 64 bins to predict the dimensions.<br>
I find I can’t make it. My loss.backward() can’t work!</p>
<pre><code class="lang-python">            output_coor_x_ = output_coor_x_.squeeze()
            output_coor_y_ = output_coor_y_.squeeze()
            output_coor_z_ = output_coor_z_.squeeze()

           ####

            output_coor_ = torch.stack([torch.argmax(output_coor_x_, axis=0),
                                     torch.argmax(output_coor_y_, axis=0),
                                     torch.argmax(output_coor_z_, axis=0)], axis=2)
            output_coor_[output_coor_ == cfg.network.coor_bin] = 0
            output_coor_ = 2.0 * output_coor_.float() / (63.0) - 1.0      # [-1,1]
            output_coor_[:, :, 0] = output_coor_[:, :, 0] * abs(model_min_x)
</code></pre>
<p>I try to use variable with before ‘####’, loss.backward() work; but when I use variable after ‘####’,  loss.backward() not work and output as follow:</p>
<p><code>RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn</code></p>
<p>I am sure that parameters’s requires_grad=True in the model.</p>
<p>Maybe I shouldn’t use torch.argmax() ?</p>

username: KFrank
cooked: <p>Hi Kiruto!</p>
<aside class="quote no-group quote-modified" data-username="KirutoCode" data-post="1" data-topic="64782" data-full="true">
<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/kirutocode/48/19029_2.png" class="avatar"> KirutoCode:</div>
<blockquote>
<p>My loss.backward() can’t work!</p>
<p>…</p>
<p><code>RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn</code></p>
<p>…</p>
<p>Maybe I shouldn’t use torch.argmax() ?</p>
</blockquote>
</aside>
<p>In short, yes, don’t use <code>argmax()</code> in your loss function.</p>
<p><code>loss.backward()</code> computes the gradient (derivative) of <code>loss</code><br>
with respect to your model parameters.  (The optimizer then uses<br>
a gradient-descent algorithm to adjust those parameters to make<br>
the loss smaller.)</p>
<p>For this to work, your loss function has to be differentiable.<br>
<code>argmax()</code> isn’t differentiable, so you can’t use it as part of<br>
your loss.</p>
<p>I don’t follow what you are doing in your code, so I don’t know<br>
specifically how you might fix it.  But as a general comment,<br>
you might investigate whether a differentiable approximation<br>
to <code>argmax()</code> would still make sense in your loss function.</p>
<p><code>softmax()</code> (which should more accurately be called <em>softargmax()</em>)<br>
can be understood to be a differentiable approximation to <code>argmax()</code><br>
(much in the same way that <code>sigmoid()</code> is a differentiable<br>
approximation to a step function).</p>
<p>If you can reformulate your loss function to use <code>softmax()</code>, you<br>
can make it differentiable and <code>loss.backward()</code> should work.</p>
<p>Good luck.</p>
<p>K. Frank</p>

username: Sourabh
cooked: <p><a class="mention" href="/u/kfrank">@KFrank</a>  <a class="mention" href="/u/ptrblck">@ptrblck</a>. Hello. I have a question in relation to using argmax.<br>
I am doing the following :</p>
<ol>
<li>A model is used as an encoder.</li>
<li>The model predicts some outputs which I then take and convert into a numpy array. (so I also detach them first from the tensor.)</li>
<li>During the course of subsequent calculations on this numpy array , I use a <code>argmax()</code> to finally return me something ( for example something like [[1,4,6,3]]. Lets call them <code>predictions</code>.)</li>
<li>I convert these predictions to tensors and also set the <code>requires_grad</code> paramter to <code>True</code> as:</li>
</ol>
<blockquote>
<p>torch.tensor(predictions, requires_grad=True)</p>
</blockquote>
<ol start="5">
<li>I then try and apply the torch.nn.L1Loss on this and my ground truth as :</li>
</ol>
<blockquote>
<p>loss=torch.nn.L1Loss(predictions,ground_truth)</p>
</blockquote>
<p>When I run the code, I see the same loss for all the epochs and the model basically does not train.<br>
For example:</p>
<blockquote>
<p>Epoch : 1, Loss : 4.0<br>
Epoch : 2, Loss : 4.0<br>
Epoch : 3, Loss : 4.0<br>
Epoch : 4, Loss : 4.0</p>
</blockquote>
<p>Now , I know that the <code>argmax()</code> is not differentiable and hence should not be used in the loss. However, I try and do some post processing on my outputs (that involves me doing the argmax() operation) before converting it to a tensor and then apply the <code>CrossEntropy</code> loss to it.<br>
If it is still wrong, it there an alternative to using the argmax()?<br>
Could you please help me with this??</p>

username: KFrank
cooked: <p>Hi Sourabh!</p>
<aside class="quote no-group" data-username="Sourabh" data-post="3" data-topic="64782" data-full="true">
<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/sourabh/48/45473_2.png" class="avatar"> Sourabh:</div>
<blockquote>
<p>I have a question</p>
</blockquote>
</aside>
<p>Please start a new thread and post your question there.  It’s preferred<br>
to resurrecting zombie threads.</p>
<p>Good luck!</p>
<p>K. Frank</p>