Loss is calculated as NaN when using amp.autocast

Anishkumar_Iyer: <p>This is my train function<br>
def _train(self):<br>
self.model.train()  # train mode<br>
train_losses = <span class="chcklst-box fa fa-square-o fa-fw"></span>  # accumulate the losses here</p>
<pre><code>        batch_iter = tqdm(
            enumerate(self.training_DataLoader),
            "Training",
            total=len(self.training_DataLoader),
            leave=False,
        )

        for i, (x, y) in batch_iter:
            input, target = x.to(self.device), y.to(
                self.device
            )  # send to device (GPU or CPU)

            self.optimizer.zero_grad()  # zerograd the parameters

            with torch.cuda.amp.autocast():
                out = self.model(input)
                # out = out.float()
                loss = self.criterion(out, target)  # calculate loss

            # print(dice_loss(target,out))
            # loss += self.criterion2(out,target)

            self.scaler.scale(loss).backward()
            self.scaler.step(self.optimizer)
            self.scaler.update()
            # torch.cuda.empty_cache()

            loss_value = loss
            train_losses.append(loss_value)

            batch_iter.set_description(
                f"Training: (loss {loss_value:.4f})"
            )  # update progressbar

        # plt.plot(plt_trainloss)
        self.training_loss.append(np.mean(train_losses))
</code></pre>
<p>I am doing a binary segmentation task where my model returns logits. when I trued to use the AMP, I get my loss as NAN. Can anyone help me figure out whats wrong here ?</p>

ptrblck: <p>Try to narrow down where the first invalid value is created by checking the loss, calculation as well as the intermediates of your forward pass.</p>

Anishkumar_Iyer: <aside class="quote no-group" data-username="Anishkumar_Iyer" data-post="1" data-topic="186272">
<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/anishkumar_iyer/48/62001_2.png" class="avatar"> Anishkumar_Iyer:</div>
<blockquote>
<pre><code class="lang-auto">with torch.cuda.amp.autocast():
                out = self.model(input)
                # out = out.float()
                loss = self.criterion(out, target)  # calculate loss
</code></pre>
</blockquote>
</aside>
<p>Hey thanks for replying ! I’m Pretty new to pytorch and deep learning in general.<br>
The NaN loss occurs at the first epoch, when the first batch is iterated<br>
I Tried to print the loss here<br>
with torch.cuda.amp.autocast():<br>
out = self.model(input)</p>
<pre><code>            loss = self.criterion(out, target)  # calculate loss
            print(loss)
</code></pre>
<p>I get this output:<br>
ensor(nan, device=‘cuda:0’, grad_fn=)<br>
The loss is calculated as nan automatically in the autocast loop before the gradients can be updated</p>

Anishkumar_Iyer: <p>I dont know the Model give NaN outputs on the first batch in autocast, It works well without autocast:</p>
<p>Is something wrong with the model ? Should I modify anything here<br>
I’m Implementing a Unet here:<br>
“”" Full assembly of the parts to form the complete network “”"<br>
import torch<br>
import torch.nn as nn<br>
import torch.nn.functional as F</p>
<p>class Unet(nn.Module):<br>
def <strong>init</strong>(self, n_channels, n_classes, bilinear=False):<br>
super(Unet, self).<strong>init</strong>()<br>
self.n_channels = n_channels<br>
self.n_classes = n_classes<br>
self.bilinear = bilinear</p>
<pre><code>    self.inc = (DoubleConv(n_channels, 64))
    self.down1 = (Down(64, 128))
    self.down2 = (Down(128, 256))
    self.down3 = (Down(256, 512))
    factor = 2 if bilinear else 1
    self.down4 = (Down(512, 1024 // factor))
    self.up1 = (Up(1024, 512 // factor, bilinear))
    self.up2 = (Up(512, 256 // factor, bilinear))
    self.up3 = (Up(256, 128 // factor, bilinear))
    self.up4 = (Up(128, 64, bilinear))
    self.outc = (OutConv(64, n_classes))

def forward(self, x):
    x1 = self.inc(x)
    x2 = self.down1(x1)
    x3 = self.down2(x2)
    x4 = self.down3(x3)
    x5 = self.down4(x4)
    x = self.up1(x5, x4)
    x = self.up2(x, x3)
    x = self.up3(x, x2)
    x = self.up4(x, x1)
    logits = self.outc(x)
    return logits
</code></pre>
<p>class DoubleConv(nn.Module):<br>
“”“(convolution =&gt; [BN] =&gt; ReLU) * 2"”"</p>
<pre><code>def __init__(self, in_channels, out_channels, mid_channels=None):
    super().__init__()
    if not mid_channels:
        mid_channels = out_channels
    self.double_conv = nn.Sequential(
        nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
        nn.BatchNorm2d(mid_channels),
        nn.ReLU(inplace=True),
        nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True)
    )

def forward(self, x):
    return self.double_conv(x)
</code></pre>
<p>class Down(nn.Module):<br>
“”“Downscaling with maxpool then double conv”“”</p>
<pre><code>def __init__(self, in_channels, out_channels):
    super().__init__()
    self.maxpool_conv = nn.Sequential(
        nn.MaxPool2d(2),
        DoubleConv(in_channels, out_channels)
    )

def forward(self, x):
    return self.maxpool_conv(x)
</code></pre>
<p>class Up(nn.Module):<br>
“”“Upscaling then double conv”“”</p>
<pre><code>def __init__(self, in_channels, out_channels, bilinear=True):
    super().__init__()

    # if bilinear, use the normal convolutions to reduce the number of channels
    if bilinear:
        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
    else:
        self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
        self.conv = DoubleConv(in_channels, out_channels)

def forward(self, x1, x2):
    x1 = self.up(x1)
    # input is CHW
    diffY = x2.size()[2] - x1.size()[2]
    diffX = x2.size()[3] - x1.size()[3]

    x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
                    diffY // 2, diffY - diffY // 2])
    # if you have padding issues, see
    # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
    # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
    x = torch.cat([x2, x1], dim=1)
    return self.conv(x)
</code></pre>
<p>class OutConv(nn.Module):<br>
def <strong>init</strong>(self, in_channels, out_channels):<br>
super(OutConv, self).<strong>init</strong>()<br>
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)</p>
<pre><code>def forward(self, x):
    return self.conv(x)
</code></pre>