Cross Entropy Loss outputting Nan

username: PCPJ
cooked: <p>I am trying to train a model with the Kitti road segmentation dataset: <a href="http://www.cvlibs.net/datasets/kitti/eval_road.php" rel="noopener nofollow ugc">http://www.cvlibs.net/datasets/kitti/eval_road.php</a>.</p>
<p>I am getting Nan from the CrossEntropyLoss module. Notice that it is returning Nan already in the first mini-batch. I already checked my input tensor for Nans and Infs.<br>
The tensor shapes I am giving to the loss func are: (b_size, n_class, h, w) and (b_size, h, w).</p>
<p>When I try to reshape the tensor in the following way:</p>
<pre><code class="lang-auto">loss = criterion(prediction.permute(0,2,3,1).contiguous().view(-1, n_class), target.view(-1))
</code></pre>
<p>the Nans disappear.<br>
Someone knows what can be happening?<br>
I am afraid that maybe something is still wrong, because I am still having troubles too make the network to converge.</p>

username: ptrblck
cooked: <p>Could you upload the <code>prediction</code> and <code>target</code> tensors so that we could try to reproduce this issue, please?<br>
Also, could you post the output of <code>python -m torch.utils.collect_env</code>?</p>

username: PCPJ
cooked: <p>Thanks for replying <a class="mention" href="/u/ptrblck">@ptrblck</a>.<br>
Here is the prediction and target tensors:<br>
<a href="https://drive.google.com/file/d/1XjKMt4kBv8Ipzm0ttWlrp906p8LXdTy1/view?usp=sharing" class="onebox" target="_blank" rel="noopener nofollow ugc">https://drive.google.com/file/d/1XjKMt4kBv8Ipzm0ttWlrp906p8LXdTy1/view?usp=sharing</a><br>
Here is the environment information:</p>
<pre><code class="lang-auto">Collecting environment information...
PyTorch version: 1.8.0+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.2 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: version 3.16.3

Python version: 3.8 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: GeForce RTX 2060
Nvidia driver version: 460.73.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.2.0
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.5
[pip3] torch==1.8.0+cu111
[pip3] torchaudio==0.8.0
[pip3] torchvision==0.9.0+cu111
[conda] Could not collect
</code></pre>

username: ptrblck
cooked: <p>Thanks for the data. It seems you are using a manual <code>float16</code> approach, which could easily overflow.<br>
If I wrap the loss calculation into <code>with torch.cuda.amp.autocast()</code> the result will be valid.<br>
What’s the use case you are not using <code>autocast</code>, but a manual approach?</p>

username: PCPJ
cooked: <p>I am training with half precision because the final hardware I will deploy my model works with half precision.<br>
I was doing this cast manually because I though it was the right way.<br>
Sorry about that.<br>
But there is still one thing without answer. Why when I changed the shape of the tensor the Nans disappeared?</p>
<p>So, to summarize, I shouldn’t train my model with half precision this way:</p>
<pre><code class="lang-auto">model = Model()
model.to('cuda').half()
for layer in model.modules():
    if isinstance(layer, torch.nn.BatchNorm2d):
        layer.float()

for i, (x y) in enumerate(trainLoader):
        x = x.to('cuda').half()
        y = y.to('cuda').half()
</code></pre>
<p>Where should I put the <code>with torch.cuda.amp.autocast()</code>?</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="PCPJ" data-post="5" data-topic="123002">
<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/p/47e85d/48.png" class="avatar"> PCPJ:</div>
<blockquote>
<p>But there is still one thing without answer. Why when I changed the shape of the tensor the Nans disappeared?</p>
</blockquote>
</aside>
<p>Depending on the shapes, different methods will be called as seen <a href="https://github.com/pytorch/pytorch/blob/4cb534f92ef6f5b2ec99109b0329f93a859ae831/aten/src/ATen/native/LossNLL.cpp#L487-L494">here</a>.<br>
In particular you could compare the methods used in <a href="https://github.com/pytorch/pytorch/blob/4cb534f92ef6f5b2ec99109b0329f93a859ae831/aten/src/ATen/native/LossNLL.cpp#L29"><code>nll_loss_out_frame</code></a> and <a href="https://github.com/pytorch/pytorch/blob/4cb534f92ef6f5b2ec99109b0329f93a859ae831/aten/src/ATen/native/LossNLL2d.cpp#L84"><code>nll_loss2d_forward_out_frame</code></a>, which might have a different order of accumulation (and the latter seems to over-/underflow in <code>float16</code>) and is why they promote to <code>float32</code> using <code>amp</code>.</p>
<aside class="quote no-group" data-username="PCPJ" data-post="5" data-topic="123002">
<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/p/47e85d/48.png" class="avatar"> PCPJ:</div>
<blockquote>
<p>Where should I put the <code>with torch.cuda.amp.autocast()</code>?</p>
</blockquote>
</aside>
<p>You can wrap the model execution in it as described in the <a href="https://pytorch.org/docs/stable/notes/amp_examples.html#amp-examples">examples</a>:</p>
<pre data-code-wrap="python"><code class="lang-python"># Creates model and optimizer in default precision
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

# Creates a GradScaler once at the beginning of training.
scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()

        # Runs the forward pass with autocasting.
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)

        # Scales loss.  Calls backward() on scaled loss to create scaled gradients.
        # Backward passes under autocast are not recommended.
        # Backward ops run in the same dtype autocast chose for corresponding forward ops.
        scaler.scale(loss).backward()

        # scaler.step() first unscales the gradients of the optimizer's assigned params.
        # If these gradients do not contain infs or NaNs, optimizer.step() is then called,
        # otherwise, optimizer.step() is skipped.
        scaler.step(optimizer)

        # Updates the scale for next iteration.
        scaler.update()
</code></pre>
<aside class="quote no-group" data-username="PCPJ" data-post="5" data-topic="123002">
<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/p/47e85d/48.png" class="avatar"> PCPJ:</div>
<blockquote>
<p>So, to summarize, I shouldn’t train my model with half precision this way:</p>
</blockquote>
</aside>
<p>You could try to manually cast the model to <code>float16</code>, but as the model can easily create invalid outputs, you would also have to verify, that all operations are working correctly.</p>