Conv2d.backwards always results in NaN

username: Tommy1
cooked: <p>I’ve only started using Torch very recently, so please forgive me if I’ve made some obvious error. I have been trying to convert an old TensorFlow notebook of mine to Torch, but I am currently running into what I see as an odd issue: the torch <code>backwards</code> function, when run on my network, always produces NaN results (thus causing the weights to be adjusted to NaN after one step of optimization). There is no issue with feeding the network forward, and from what I can tell from stepping through the process manually, the NaN weights solely spawn from the Conv2d network. Any help on figuring out what the problem is would be greatly appreciated.</p>
<p>Because it’s practically impossible to debug something by description alone, here’s a tiny reproduction of the problematic section of my code.</p>
<pre><code class="lang-auto">device = "cuda" # Fails for both "cuda" and "cpu" on a Colab environment

binary_crossentropy = nn.BCELoss()

def discriminator_loss(real_guess, fake_guess):
  loss_for_target = binary_crossentropy(torch.ones_like(real_guess), real_guess)
  loss_for_predicted = binary_crossentropy(torch.zeros_like(fake_guess), fake_guess)

  return loss_for_target + loss_for_predicted

with torch.autograd.detect_anomaly():
  class MiniModel(nn.Module):
    def __init__(self):
      super(MiniModel, self).__init__()
      self.last = nn.Conv2d(512, 1, 4, 1)
    def forward(self, x):
      return self.last(x)
  
  mm = MiniModel().to(device)
  opt = optim.Adam(mm.parameters())
  inp1 = torch.rand((16, 512, 14, 14)).to(device)
  inp2 = torch.rand((16, 512, 14, 14)).to(device)
  
  out = discriminator_loss(mm.last(inp1), mm.last(inp2))
  print(out)

  grads = out.backward()
  opt.step()
  print(list(mm.parameters()))
</code></pre>
<p>The results:</p>
<pre><code class="lang-auto">tensor(100.1079, device='cuda:0', grad_fn=&lt;AddBackward0&gt;)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-52-e77288483237&gt; in &lt;module&gt;()
     25   print(out)
     26 
---&gt; 27   grads = out.backward()
     28   opt.step()
     29   print(list(mm.parameters()))

/usr/local/lib/python3.7/dist-packages/torch/_tensor.py in backward(self, gradient, retain_graph, create_graph, inputs)
    305                 create_graph=create_graph,
    306                 inputs=inputs)
--&gt; 307         torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
    308 
    309     def register_hook(self, hook):

/usr/local/lib/python3.7/dist-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)
    154     Variable._execution_engine.run_backward(
    155         tensors, grad_tensors_, retain_graph, create_graph, inputs,
--&gt; 156         allow_unreachable=True, accumulate_grad=True)  # allow_unreachable flag
    157 
    158 

RuntimeError: Function 'CudnnConvolutionBackward0' returned nan values in its 1th output.
</code></pre>
<p>Interestingly (at least for me), when I use the <code>cpu</code> instead of the <code>cuda</code> device, the <code>detect_anomaly()</code> does not trigger a RuntimeError, but the weights are still readjusted to NaNs:</p>
<pre><code class="lang-auto">[Parameter containing:
tensor([[[[nan, nan, nan, nan],
          [nan, nan, nan, nan],
          [nan, nan, nan, nan],
...
</code></pre>
<p>As a final note, I <em>have</em> verified that the weights are not initially NaNs; I just didn’t think it necessary to include the code showing that here.</p>
<p>Again, any help would be greatly appreciated; I’m really liking Torch’s API so far, but I perhaps I should have started on a simpler model.</p>

username: Andrei_Cristea
cooked: <p>Try making the following changes:</p>
<p>inside your loss fn:</p>
<pre><code class="lang-auto">  loss_for_target = binary_crossentropy(real_guess, torch.ones_like(real_guess).to(device))  # !! invert order of input and target
  loss_for_predicted = binary_crossentropy(fake_guess, torch.zeros_like(fake_guess).to(device))  # !! invert order of input and target
</code></pre>
<p>inside your model, make sure you return values between 0 and 1 (else we can’t compute BCELoss). something like this:<br>
<code>   return nn.Sigmoid()(self.last(x))  # !! force values to be in (0, 1) interval</code></p>
<p>finally, make sure you apply your entire model, not just one layer:<br>
<code>  out = discriminator_loss(mm(inp1), mm(inp2))  # !! applying the entire model, not just the "last"</code></p>
<p>With these changes, the code runs for me.</p>

username: ptrblck
cooked: <p>You are passing the arguments to the criterion in the wrong order and should also use <code>nn.BCEWithLogitsLoss</code> based on the model output.</p>
<p><a class="mention" href="/u/andrei_cristea">@Andrei_Cristea</a> was faster! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

username: Tommy1
cooked: <p>Ah, thank you both <a class="mention" href="/u/andrei_cristea">@Andrei_Cristea</a> and <a class="mention" href="/u/ptrblck">@ptrblck</a>! This was definitely an issue of converting from TensorFlow without fully understanding the differences; TF has a <code>from_logits</code> argument in its BinaryCrossentropy class, while Torch provides two separate classes. Similarly, they swap the order of true and false labels when applying the loss function. It seems that I should read the documentation more carefully rather than assuming that similarly named classes have similar calling conventions across the two APIs.</p>