ZeroDivisionError and Loss goes to NaN with Apex Loss Scaling

username: Arjun_Gupta
cooked: <p>I am using an open source distributed PyTorch implementation of training AlexNet from scratch on ImageNet (<a href="https://github.com/richardkxu/distributed-pytorch" rel="nofollow noopener">https://github.com/richardkxu/distributed-pytorch</a>).</p>
<p>This implementation works flawlessly as is. As soon as I add an additional loss (loss_contrastive) in the following manner:</p>
<pre><code class="lang-auto">loss = criterion(output, target)
loss_contrastive = getContrastiveLoss(target, rep3, rep4, rep5, contrastive_idxs)
loss += 0.1*loss_contrastive

optimizer.zero_grad()

# Mixed-precision training requires that the loss is scaled in order
# to prevent the gradients from underflow
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()
</code></pre>
<p>I get a ZeroDivisionError on the last line. Also, I am getting a gradient overflow error for many consecutive steps (<code>Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 5e-324</code>) and looking at the two losses, both losses separately start at around ~10, and then <code>loss_contastive</code> begins rapidly increasing. After many steps of <code>loss_contrastive</code> being at around ~10^8 and many gradient overflows (here the original loss is ~50), both losses become NaNs.</p>
<p>For context, <code>loss_contrastive</code> is simply a contrastive MSE loss which aims to minimize the distance between certain representations for certain inputs and maximize it for others. Am I treating the addition of a new loss incorrectly? Any ideas what might be causing this?</p>
<p>Thanks!</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="Arjun_Gupta" data-post="1" data-topic="91637">
<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/arjun_gupta/48/12756_2.png" class="avatar"> Arjun_Gupta:</div>
<blockquote>
<p>After many steps of <code>loss_contrastive</code> being at around ~10^8 and many gradient overflows (here the original loss is ~50), both losses become NaNs.</p>
</blockquote>
</aside>
<p>This doesn’t seem right. Are you expecting the loss to have such a high value?<br>
If the loss gets NaN values, could you check, if your model is also outputting NaNs?</p>
<p>The zero division error is raised, as the loss scaler detects NaN gradients (due to an overflow, which might sometimes happen and is expected, or due to the NaN loss, which should never happen), and thus reduces the scaling factor until it’s not representable anymore.</p>
<p>That being said, we recommend to use the <a href="https://pytorch.org/docs/stable/amp.html">native amp implementation</a> instead of <code>apex</code>.</p>

username: Arjun_Gupta
cooked: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>,</p>
<p>Thanks so much for your response. I decreased the learning rate by an order of magnitude as per someone’s suggestion on a related GitHub Issue, and this seemed to resolve the problem! And thank you for your suggestion – I will take a look into the native amp implementation.</p>
<p>I do have a quick follow-up though: Given that the original loss trained well with a learning rate of <code>0.01</code>, and the new loss trains well with a different learning rate of <code>0.001</code>, I thought it would make sense to use two separate optimizers in the following manner:</p>
<pre><code class="lang-auto">optimizer = torch.optim.SGD(model.parameters(), 
                            0.01, 
                            momentum=args.momentum, 
                            weight_decay=args.weight_decay)
parameters_contrastive = list(model.conv2.parameters()) + 
                         list(model.maxpool2.parameters()) + 
                         list(model.conv3.parameters())
optimizer_contrastive = torch.optim.SGD(parameters_contrastive, 
                                        0.001,
                                        momentum=args.momentum,
                                        weight_decay=args.weight_decay)

model, [optimizer, optimizer_contrastive] = amp.initialize(model, 
                                                           [optimizer, optimizer_contrastive],
                                                           opt_level=args.opt_level,
                                                           keep_batchnorm_fp32=args.keep_batchnorm_fp32,
                                                           loss_scale=args.loss_scale,
                                                           num_losses=2)
</code></pre>
<p>and in the train loop:</p>
<pre><code class="lang-auto">output, rep3, rep4, rep5 = model(input)

loss = criterion(output, target)

# compute contrastive loss
loss_contrastive = getContrastiveLoss(target, rep3, rep4, rep5, contrastive_idxs)
loss_contrastive *= 0.1

# compute gradient and do SGD step
optimizer.zero_grad()
optimizer_contrastive.zero_grad()

# Mixed-precision training requires that the loss is scaled in order
# to prevent the gradients from underflow
with amp.scale_loss(loss, optimizer, loss_id=0) as scaled_loss:
    scaled_loss.backward(retain_graph=True)
with amp.scale_loss(loss_contrastive, optimizer_contrastive, loss_id=1) as scaled_loss:
    scaled_loss.backward()

optimizer.step()
optimizer_contrastive.step()
</code></pre>
<p>While both losses (<code>loss</code> and <code>loss_contrastive</code>) start off at the same order of magnitude (after <code>loss_contrastive</code> is scaled by 0.1), <code>loss_contrastive</code> is reduced very smoothly throughout training while <code>loss</code> remains pretty much constant. Again, just to reiterate from my original post, when this exact model is trained without the additional <code>loss_contrastive</code>, the <code>loss</code> is reduced very smoothly and the AlexNet is able to achieve a final top-5 accuracy of 80% on ImageNet.</p>
<p>I am unsure how the addition of my <code>loss_contrastive</code> is discouraging the network from optimizing the original, main <code>loss</code>. Do you have any idea why this might be the case, or what I can do to combat this?</p>
<p>Thanks once again for your help; You have helped me numerous times in the past as well and it always stuns me to see how people of your calibre are so generous and accessible even to beginners. I truly appreciate all your effort.</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="Arjun_Gupta" data-post="3" data-topic="91637">
<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/arjun_gupta/48/12756_2.png" class="avatar"> Arjun_Gupta:</div>
<blockquote>
<p>I am unsure how the addition of my <code>loss_contrastive</code> is discouraging the network from optimizing the original, main <code>loss</code> . Do you have any idea why this might be the case, or what I can do to combat this?</p>
</blockquote>
</aside>
<p>This is pure speculation, but I think that in your use case where both losses are used, the gradients of <code>loss_contrastive</code> might not yield a proper signal and in fact your model mainly optimizes for the <code>loss</code>.<br>
This might happen, if the gradient magnitude is different, so that the gradients of <code>loss_contrastive</code> could be seen as noise compared to the gradients of <code>loss</code>.<br>
If that’s the case, the optimization might lower one loss and unfortunately increase the other one.<br>
You could try to use different weightings on the losses directly to create an equal range of their values and rerun the code again.</p>
<aside class="quote no-group" data-username="Arjun_Gupta" data-post="3" data-topic="91637">
<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/arjun_gupta/48/12756_2.png" class="avatar"> Arjun_Gupta:</div>
<blockquote>
<p>Thanks once again for your help; You have helped me numerous times in the past as well and it always stuns me to see how people of your calibre are so generous and accessible even to beginners. I truly appreciate all your effort.</p>
</blockquote>
</aside>
<p>Thanks for the kind words and I’m glad I could help. <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>