Backpropagation issue when using argmax

username: blade
cooked: <p>I have a network with following architecture</p>
<pre><code>def forward(self, W, z):

    W1 = self.SM(W)
    W2 = torch.argmax(W1, dim=3).float()

    h_5 = W2.view(-1, self.n_max_atom * self.n_max_atom)
    h_6 = self.leaky((self.dec_fc_5(h_5)))
    h_6 = h_6.view(-1, self.n_max_atom, self.n_atom_features)

    return h_6
</code></pre>
<p>where  <code>self.SM = nn.Softmax(dim=3)</code>. When running this, I receive the error:</p>
<blockquote>
<p>File “/Users/Blade/model/VAEtrain.py”, line 154, in trainepoch<br>
loss.backward()<br>
File “/Users/Blade/anaconda3/lib/python3.7/site-packages/torch/tensor.py”, line 198, in backward<br>
torch.autograd.backward(self, gradient, retain_graph, create_graph)<br>
File “/Users/Blade/anaconda3/lib/python3.7/site-packages/torch/autograd/<strong>init</strong>.py”, line 100, in backward<br>
allow_unreachable=True)  # allow_unreachable flag<br>
RuntimeError: Expected isFloatingType(grads[i].scalar_type()) to be true, but got false.  (Could this error message be improved?  If so, please report an enhancement request to PyTorch.) (validate_outputs at /Users/distiller/project/conda/conda-bld/pytorch_1587428061935/work/torch/csrc/autograd/engine.cpp:476)<br>
frame <span class="hashtag-raw">#0:</span> c10::Error::Error(c10::SourceLocation, std::__1::basic_string&lt;char, std::__1::char_traits, std::__1::allocator &gt; const&amp;) + 135 (0x112602ab7 in libc10.dylib)<br>
frame <span class="hashtag-raw">#1:</span> torch::autograd::validate_outputs(std::__1::vector&lt;torch::autograd::Edge, std::__1::allocator<a>torch::autograd::Edge</a> &gt; const&amp;, std::__1::vector&lt;at::Tensor, std::__1::allocator<a>at::Tensor</a> &gt;&amp;, std::__1::function&lt;std::__1::basic_string&lt;char, std::__1::char_traits, std::__1::allocator &gt; (std::__1::basic_string&lt;char, std::__1::char_traits, std::__1::allocator &gt; const&amp;)&gt; const&amp;) + 5884 (0x1193d614c in libtorch_cpu.dylib)<br>
frame <span class="hashtag-raw">#2:</span> torch::autograd::Engine::evaluate_function(std::__1::shared_ptr<a>torch::autograd::GraphTask</a>&amp;, torch::autograd::Node*, torch::autograd::InputBuffer&amp;) + 1996 (0x1193d18ec in libtorch_cpu.dylib)<br>
frame <span class="hashtag-raw">#3:</span> torch::autograd::Engine::thread_main(std::__1::shared_ptr<a>torch::autograd::GraphTask</a> const&amp;, bool) + 497 (0x1193d08b1 in libtorch_cpu.dylib)<br>
frame <span class="hashtag-raw">#4:</span> torch::autograd::Engine::thread_init(int) + 152 (0x1193d0648 in libtorch_cpu.dylib)<br>
frame <span class="hashtag-raw">#5:</span> torch::autograd::python::PythonEngine::thread_init(int) + 52 (0x111b43a04 in libtorch_python.dylib)<br>
frame <span class="hashtag-raw">#6:</span> void* std::__1::__thread_proxy&lt;std::__1::tuple&lt;std::__1::unique_ptr&lt;std::__1::__thread_struct, std::__1::default_delete<a>std::__1::__thread_struct</a> &gt;, void (torch::autograd::Engine::<em>)(int), torch::autograd::Engine</em>, int&gt; &gt;(void*) + 66 (0x1193dfc32 in libtorch_cpu.dylib)<br>
frame <span class="hashtag-raw">#7:</span> _pthread_body + 126 (0x7fff6e8632eb in libsystem_pthread.dylib)<br>
frame <span class="hashtag-raw">#8:</span> _pthread_start + 66 (0x7fff6e866249 in libsystem_pthread.dylib)<br>
frame <span class="hashtag-raw">#9:</span> thread_start + 13 (0x7fff6e86240d in libsystem_pthread.dylib)</p>
</blockquote>
<p>It seems that the torch.argmax function brakes the backpropagation. But I believe that it should work fine. Does the model architecture have a problem?</p>

username: KFrank
cooked: <p>Hello Blade!</p>
<aside class="quote no-group" data-username="blade" data-post="1" data-topic="100335" 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/blade/48/23861_2.png" class="avatar"> blade:</div>
<blockquote>
<p>It seems that the torch.argmax function brakes the backpropagation.</p>
</blockquote>
</aside>
<p><code>argmax()</code> is not usefully differentiable, and so, indeed, does break<br>
backpropagation.</p>
<aside class="quote no-group">
<blockquote>
<p>Does the model architecture have a problem?</p>
</blockquote>
</aside>
<p>Yes, the <code>argmax()</code> piece isn’t differentiable.</p>
<p>Make sure you understand why <code>argmax()</code> isn’t differentiable, and then<br>
see if you can reformulate what you are doing in a way that avoids the<br>
discrete jumps inherent in <code>argmax()</code>.</p>
<p>Good luck.</p>
<p>K. Frank</p>

username: blade
cooked: <p>Thanks! Is there a work around for this?</p>

username: KFrank
cooked: <p>Hello Blade!</p>
<aside class="quote no-group" data-username="blade" data-post="3" data-topic="100335" 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/blade/48/23861_2.png" class="avatar"> blade:</div>
<blockquote>
<p>Is there a work around for this?</p>
</blockquote>
</aside>
<p>“Work around” isn’t really the right term, as it implies that you are<br>
trying to do something that makes sense, and you need to work<br>
around a bug or limitation to reach the same result by a different<br>
path.</p>
<p>Conceptually, what are the index values in <code>W2</code> supposed to mean?<br>
How do you want your optimizer to respond when one of the values<br>
in <code>W2</code> suddenly makes a discrete jump from <code>2</code> to <code>1</code>?</p>
<p>To illustrate my point, here is a “work around”:</p>
<p>Replace <code>W2 = torch.argmax (W1, dim = 3).float()</code><br>
with <code>W2 = 0.0 * torch.sum (W1, dim = 3).float()</code>.</p>
<p>Backpropagation will now work (but all of your gradients will be zero).</p>
<p><code>softmax()</code> is a smooth (differentiable) approximation to the one-hot<br>
encoding of <code>argmax()</code>.  But this comment will only be helpful if you<br>
understand the conceptual role you want the piece-wise constant (not<br>
usefully differentiable) <code>W2</code> values to play in your network training.</p>
<p>Good luck.</p>
<p>K. Frank</p>

username: blade
cooked: <aside class="quote no-group" data-username="KFrank" data-post="4" data-topic="100335">
<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/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>Conceptually, what are the index values in <code>W2</code> supposed to mean?</p>
</blockquote>
</aside>
<p>I have two networks that are trained together. The networks are in series, i.e. output of one is the input of the other. Upstream network is solving a classification problem for graph edge types <code>[Batch, node, node, class]</code>. What you see is my downstream network: it takes in output of the upstream <code>W</code>, applies Softmax to turn them into probabilities <code>W1</code>, and then I want to turn them into either a one-hot vector or class labels <code>W2</code> before reshaping and feeding it to the linear layer <code>dec_fc_5</code>.</p>

username: KFrank
cooked: <p>Hello Blade!</p>
<aside class="quote no-group quote-modified" data-username="blade" data-post="5" data-topic="100335" 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/blade/48/23861_2.png" class="avatar"> blade:</div>
<blockquote>
<aside class="quote no-group" data-username="KFrank" data-post="4" data-topic="100335">
<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/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>Conceptually, what are the index values in <code>W2</code> supposed to mean?</p>
</blockquote>
</aside>
<p>I have two networks that are trained together.<br>
…<br>
What you see is my downstream network: it takes in output of the upstream <code>W</code><br>
…<br>
I want to turn them into either a one-hot vector or class labels <code>W2</code></p>
</blockquote>
</aside>
<p>Once you turn the output of your upstream network into a one-hot<br>
vector or a class label, you have done something that is not<br>
differentiable and that breaks backpropagation.</p>
<p>So …   You have to do something else.</p>
<p>Why not just pass the output of your upstream directly to your<br>
downstream network (i.e., directly to <code>fc_5</code>)?  What breaks?</p>
<p>If the output of your upstream network comes directly from a <code>Linear</code><br>
without any subsequent activation function, you will want an activation<br>
function in between.  How about <code>relu()</code> or <code>sigmoid()</code> (or even<br>
<code>softmax()</code>)?</p>
<p>If passing a one-hot vector to your downstream vector makes the<br>
most sense (ignoring the fact that it isn’t differentiable), perhaps you<br>
should consider my observation that <code>softmax()</code> is a differentiable<br>
approximation to the one-hot encoding of <code>argmax()</code>.</p>
<p>You can sharpen this result – making <code>softmax()</code> “less soft” – by<br>
scaling its argument, e.g., <code>softmax (scale * W1)</code>.  As <code>scale</code> is<br>
increased to approach <code>+inf</code>, the result of <code>softmax (scale * W1)</code><br>
will approach <code>one_hot (argmax (W1))</code>.</p>
<p>Best.</p>
<p>K. Frank</p>

username: blade
cooked: <p>That makes sense, thank you very much for your thorough explanation!</p>