Unexplained SGD behaviour

kirk86: <p>Can someone please shed some light on what’s going on here? Why do we get exploding loss leading to nan?</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/2/2318d7e9623e784883a2bfba85cb100ed5d62e6b.png" data-download-href="https://discuss.pytorch.org/uploads/default/2318d7e9623e784883a2bfba85cb100ed5d62e6b" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/2/2318d7e9623e784883a2bfba85cb100ed5d62e6b_2_119x500.png" alt="image" data-base62-sha1="50tWauXGjJD0UGA0AMLep4cazjl" width="119" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/2/2318d7e9623e784883a2bfba85cb100ed5d62e6b_2_119x500.png, https://discuss.pytorch.org/uploads/default/optimized/2X/2/2318d7e9623e784883a2bfba85cb100ed5d62e6b_2_178x750.png 1.5x, https://discuss.pytorch.org/uploads/default/original/2X/2/2318d7e9623e784883a2bfba85cb100ed5d62e6b.png 2x" data-dominant-color="292E29"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">234×982 64.8 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>MWE:</p>
<pre><code class="lang-auto">import torch
import numpy as np
import torchvision as tv
import torch.nn.functional as F

stats = ((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))

train_loader = torch.utils.data.DataLoader(                                                                                                                                                  
    tv.datasets.CIFAR10(root='./data', train=False,
    transform=tv.transforms.Compose([                                                                                                                                    
                            tv.transforms.Resize(64),                                                                                                                   tv.transforms.ToTensor(),                                                                                                                   tv.transforms.Normalize(*stats)                                                                                                                                  
                        ]), download=True),                                                                                                                                                  
    batch_size=512, shuffle=True,                                                                                                                                                            
    num_workers=4, pin_memory=False)

dataiter = next(iter(train_loader))
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = tv.models.AlexNet(num_classes=10).to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, weight_decay=5e-4, momentum=0.9, nesterov=True)

for epoch in range(150):
    model.train()                                                                                                                                                                            
    for data, target in [dataiter] * 50:
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        outputs = model(data)
        outputs = F.log_softmax(outputs, dim=1)
        loss = F.nll_loss(outputs, target)
        print("loss: {}".format(loss.item()))
        probs, preds = torch.max(outputs, 1)
        loss.backward()
        optimizer.step()
</code></pre>

Zhihao_Wu: <p>When did you meet this nan? I cannot reproduce it…</p>

kirk86: <p>Seriously you can’t? So, for you the MWE works fine you don’t have any of those exploding loss issues? WTF? The nan came trying the above example and I think it’s super sensitive to the learning rate if you play a bit with it you might hit the hot spot sooner than later. I’m saying this cause I first met this nan when using a step scheduler and randomly everything started to explode. I was looking the same 10 lines of training code over and over again questioning myself whether the optimizer.zero_grad() should go before the loss or after, and all sorts of these crazy stuff.</p>

andreaskoepf: <p><a class="mention" href="/u/kirk86">@kirk86</a> I guess the learning rate is slightly too high. Do you still get NaNs with a lr of 0.01? You could test whether gradually warming up the learning rate is beneficial.</p>

andreaskoepf: <p><a class="mention" href="/u/kirk86">@kirk86</a> I played a bit with your script and indeed I can confirm that it is unexpected hard to get at quick convergence to near 0 loss for the single-test-batch. It immediately works for me with a significantly smaller learning rate like 1e-3 and small batch-sizes like 1,4 or 32  (no time to test others).</p>

kirk86: <p><a class="mention" href="/u/andreaskoepf">@andreaskoepf</a> Thanks for the reply, yup with a lr of 0.01 it works fine but what throws me off is this insane sensitivity of the optimizer. If I switch to adam I don’t get nan but I get stagnation so loss is ~same.</p>
<aside class="quote no-group" data-username="andreaskoepf" data-post="4" data-topic="54608">
<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/andreaskoepf/48/15327_2.png" class="avatar"> andreaskoepf:</div>
<blockquote>
<p>You could test whether gradually warming up the learning rate is beneficial.</p>
</blockquote>
</aside>
<p>What do we mean these days by warming up the learning rate, there are so many heuristics that’s its hard to keep track of them all.</p>

kirk86: <p><a class="mention" href="/u/andreaskoepf">@andreaskoepf</a> Appreciate it! But it’s strange cause this way with smaller batch sizes and learning rates we loose all the benefits of gpus and fast training, so in a sense to me it seems that we’re back at square one.</p>

SimonW: <p>On GPU it is normal to tune the lr. Too large lr can cause explosion.</p>

andreaskoepf: <aside class="quote no-group" data-username="kirk86" data-post="6" data-topic="54608">
<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/kirk86/48/151_2.png" class="avatar"> kirk86:</div>
<blockquote>
<p>What do we mean these days by warming up the learning rate,</p>
</blockquote>
</aside>
<p>I thought of gradually increasing the learning rate during a ramp-up period at the beginning of the training, e.g. with a learning rate scheduler like <a href="https://github.com/ildoonet/pytorch-gradual-warmup-lr" class="inline-onebox" rel="noopener nofollow ugc">GitHub - ildoonet/pytorch-gradual-warmup-lr: Gradually-Warmup Learning Rate Scheduler for PyTorch</a></p>
<p>In general I would recommend to experiment with different optimizers, weight-initializations, activation-functions and learning rates. You collect loss-logs of the experiments and plot them together to see what works best. To further analyze things you could print out the max-abs values (or a norm) of the gradient.</p>

kirk86: <p>Thanks Andreas,</p>
<aside class="quote no-group" data-username="andreaskoepf" data-post="9" data-topic="54608">
<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/andreaskoepf/48/15327_2.png" class="avatar"> andreaskoepf:</div>
<blockquote>
<p>you could print out the max-abs values (or a norm) of the</p>
</blockquote>
</aside>
<p>If you don’t mind me asking does that mean taking the max abs val of the norm of all variables that have a <code>requires_grad=True</code>? Or just specific variables/layers with grads?</p>