FP16 gives NaN loss when using pre-trained model

username: bluesky314
cooked: <p>I tried the new fp16 in native torch. However, when I continue my model training for my segmentation task I get loss as NaNs. With the same script, if I initialize the same model architecture from scratch then it works fine.<br>
I set adam’s ‘ep’ to 1e-4 as well but it made no difference. I get NaN loss from the first batch continuing my trained model.</p>
<p>The model that I am continuing is already quite good at the task. I am guessing it is due to very low loss value but am not sure. How can I fix this?</p>
<p>Python 3.6<br>
Pytorch 1.6<br>
Linux<br>
pip installed lightning<br>
CUDA Version: 10.1<br>
Nvidia-V100</p>

username: a_d
cooked: <p>Could you please share your training loop?<br>
Also  basic autocast usage would be something like -</p>
<pre><code class="lang-auto">with autocast():
    out = model(x)
    loss = criterion(targets, out)
print(loss) #to debug
loss.backward()
optim.step()
</code></pre>

username: bluesky314
cooked: <p>Sure here it is. As mentioned my model works when I train a new model from scratch. It is only when continuing my good model that i get the NaNs</p>
<pre><code class="lang-auto">from torch.cuda import amp
def mini_trainfp16(model,opt,scheduler,epochs,dataloaders,image_tracker=None,saver=None,saver2=None):
    print('------------------------------------------ Using fp16')
    scaler = amp.GradScaler()
    for epoch in range(epochs):
        for e in dataloaders.keys():
            model.train() if e=='train' else model.eval()
            curr_loader = tqdm(dataloaders[e], total=int(len(dataloaders[e])))
            for idx,data in enumerate(curr_loader):
                x = data['image'].float().to(device);labels_v=data['label'].double().to(device) 
                with amp.autocast():
                    out = model(x)
                    loss = muti_bce_loss_fusionfp16(out, labels_v) 
                
                if e=='train':
                    scaler.scale(loss).backward()
                    scaler.step(opt)
                    scaler.update()
                    opt.zero_grad()
                    scheduler.step()
</code></pre>

username: a_d
cooked: <p>You’re scaling the losses again after the autocast region, which is redundant and might lead to wrong loss I believe.<br>
Remove the <code>scaler.scale(loss)</code> and simply run <code>loss.backward()</code> outside the autocast, that might fix it…<br>
A lot of training will lead to overfitting and negligible loss, but I don’t think it should result in Nan losses</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="bluesky314" data-post="1" data-topic="94133">
<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/bluesky314/48/16527_2.png" class="avatar"> bluesky314:</div>
<blockquote>
<p>I get NaN loss from the first batch continuing my trained model.</p>
</blockquote>
</aside>
<p>This would point towards an overflow in some intermediate activations or are you also seeing a NaN output if you train the model from scratch?</p>
<aside class="quote no-group" data-username="a_d" data-post="4" data-topic="94133">
<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/a/e19b73/48.png" class="avatar"> a_d:</div>
<blockquote>
<p>Remove the <code>scaler.scale(loss)</code> and simply run <code>loss.backward()</code> outside the autocast, that might fix it…</p>
</blockquote>
</aside>
<p>No, the <code>scaler</code> is an essential part of mixed-precision training and should not be removed.<br>
<a href="https://pytorch.org/docs/stable/notes/amp_examples.html#typical-mixed-precision-training">Here</a> you can see the typical training loop with comments for each step.</p>

username: a_d
cooked: <p>Oh completely missed it, I have been training mixed precisons models wrong all along, thank you <a class="mention" href="/u/ptrblck">@ptrblck</a> for correcting and sorry for posting wrong information…</p>

username: ptrblck
cooked: <p>Happy to help! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"><br>
Which models were you training and was the mixed-precision training successful without the gradient scaling?</p>

username: a_d
cooked: <p>Mainly object detection and segmentation, specifically the fasterrcnn_resnet50_fpn from the torchvision models, though the results were not that bad either, but yes I believe it did start becoming Nan loss later down the training</p>

username: bluesky314
cooked: <p>Yes you are right, I am doing deep supervision. I have 5 outputs and some of the outputs are -inf, NaNs and some real values. The final loss is a sum of all losses so that too goes to NaN. How can I fix this? This only happens when continuing the training of my good model.</p>
<p>I am suprised that all are not NaNs. There are some valid outputs. I was suspecting it is due to a batchnorm layer having too low std but am unsure now. Any suggestions to debug?</p>

username: ptrblck
cooked: <p>You could register <a href="https://discuss.pytorch.org/t/how-can-l-load-my-best-model-as-a-feature-extractor-evaluator/17254/6">forward hooks</a> for each module and check the output. Once you could isolate the first layer creating the invalid values, you could check the input as well as the parameters to further isolate it.</p>