Basic VGG16 with autocast Mixed Precision creates NAN gradients

RupertP: <p>Hello,</p>
<p><em>full code and link to Google Colab below.</em></p>
<p>I want to use a basic VGG 16 as a feature extractor. I use VGG 16 from <code>torchvision.models</code> and remove the FC layers and the Average Pooling layer.</p>
<p>Then I create a dummy input and target and use MSE loss. When enabling <code>cuda.amp.autocast</code> some of the gradients are immediatly either infinite or NAN. Disabling <code>cuda.amp.autocast</code> works fine and does not cause any NANs.</p>
<p>I noticed that I have the same problem with my ResNet50 but after 5 iterations the NANs  in the gradients disappear. I assume this is because of</p>
<pre><code class="lang-auto">    scaler.step(optimizer)
    scaler.update()
</code></pre>
<p>But the NaNs don’t seem to disappear with the VGG16 model and after a while I get a NaN loss.</p>
<p>Here is the <a href="https://colab.research.google.com/drive/1jDB_pl4ioUXtV7Ou473-HoJT0-lZKN9n?usp=sharing" rel="nofollow noopener">Google Colab link</a> . You can set <code>use_mixed_precision</code> to either <code>True</code> or <code>False</code> in order to compare the results. Dont forget to set your runtime type to GPU.</p>
<p>Is this a problem with the VGG16 from <code>torchvision.models</code> or am I doing something wrong?</p>
<p>Full Code:</p>
<pre><code class="lang-auto">import torch
import torchvision.models as models
import torch.nn as nn

class VggExtractor(torch.nn.Module):
  def __init__(self):
    super(VggExtractor, self).__init__()
    self.vgg16 = models.vgg16(pretrained=False)
    self.feature_extractor_layers = list(self.vgg16.children())[:-2]                                  #remove FC layers and AVG pool
    self.stem = torch.nn.Sequential(*(list(self.feature_extractor_layers[0].children())[0:5]))
    self.layer1 = torch.nn.Sequential(*(list(self.feature_extractor_layers[0].children())[5:10]))
    self.layer2 = torch.nn.Sequential(*(list(self.feature_extractor_layers[0].children())[10:17]))
    self.layer3 = torch.nn.Sequential(*(list(self.feature_extractor_layers[0].children())[17:24]))
    self.layer4 = torch.nn.Sequential(*(list(self.feature_extractor_layers[0].children())[24:-1]))    # remove last maxpool
  def forward(self,x):
    x = self.stem(x)
    x = self.layer1(x)
    x = self.layer2(x)
    x = self.layer3(x)
    x = self.layer4(x)
    return x

def feedforward_and_loss(net, images, target):
  output = net(images)
  loss_function = nn.MSELoss(reduction="sum")
  loss = loss_function(output, target)
  return loss

def print_nan_found(net, i):
  nan_found=False
  for name, parameter in net.named_parameters():
    if nan_found:
      break
    if parameter.grad is not None:
      nan_found = torch.isnan(parameter.grad).any()
  print("iteration", i, " nan found ", nan_found.item())

net = VggExtractor()
net = net.cuda()
scaler = torch.cuda.amp.GradScaler()

use_mixed_precision = True # Set this to False in order to disable mixed precision
scaler = torch.cuda.amp.GradScaler()
torch.cuda.amp.autocast(enabled=use_mixed_precision)
optimizer = torch.optim.SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=0.005)

for i in range(1000):
  optimizer.zero_grad()

  images = torch.ones(1,3,512,1760)
  target = torch.zeros(1,512,32,110)
  images = images.cuda()
  target = target.cuda()


  if use_mixed_precision: 
    with torch.cuda.amp.autocast():
      loss = feedforward_and_loss(net, images, target)
    scaled_loss = scaler.scale(loss)
    scaled_loss.backward()
    scaler.step(optimizer)
    scaler.update()
    print_nan_found(net, i)
  else:
      loss = feedforward_and_loss(net, images, target)
      loss.backward()
      optimizer.step()
      print_nan_found(net)


</code></pre>
<p>Greetings Rupert</p>

ptrblck: <p>Thanks for the code.<br>
If I set <code>use_mixed_precision=False</code>, the code also outputs NaNs, so it seems your training blows up:</p>
<pre><code class="lang-python">iteration 0  nan found  False
Loss is  10132.611328125
iteration 1  nan found  True
Loss is  inf
iteration 2  nan found  True
Loss is  nan
iteration 3  nan found  True
Loss is  nan
iteration 4  nan found  True
Loss is  nan
iteration 5  nan found  True
Loss is  nan
iteration 6  nan found  True
Loss is  nan
iteration 7  nan found  True
Loss is  nan
iteration 8  nan found  True
Loss is  nan
</code></pre>
<p>I guess the learning rate is just too high (for AMP and vanilla FP32 training).<br>
If you use <code>reduction='mean'</code> or lower the learning rate, the model seems to work fine in both modes (for at least 60 epochs, as I’ve stopped it afterwards).</p>