pytorch训练时前一次迭代的loss正常后一次迭代却报nan

问题描述：训练一个有四个阶段模型的网络（可以简单理解每一阶段就是一个卷积层），每阶段输出都会进行loss计算。现将pytorch原始的ce loss改为focal loss后，网络训练了数个迭代后loss 报nan。输入数据检查过没有问题，报nan时的前一个迭代的loss是正常的。报nan的当前迭代，第一个阶段所有的卷积参数都已经是nan了。

一、问题排除

1. 因为查看过数据，完全没有问题，排除输入
2. 因为参数在报nan时已经是nan了，不是前向问题
3. 前一次迭代loss完全没问题，可能不是bp的问题，但是还没有排除是bp的问题，虽然前一次迭代loss是正常的

二、问题定位

将所有反向梯度打印出来，发现第一个阶段的所有参数梯度为nan，其它正常。

for name, param in net.named_parameters():
    print('name:{} param grad:{} param requires_grad:{}'.format(name, param.grad, param.requires_grad))

采用toch.autograd.detect_anomaly()发现loss报错为“RuntimeError: Function 'LogBackward' returned nan values in its 0th output”

with autograd.detect_anomaly():
    loss.backward()

说明是在第一阶段计算focalloss时，bp出现了nan。

三、问题发生原因

3.1 为什么bp出现了nan，但是forward计算loss没有inf之类的东西出现？

因为focal loss计算的时候采用了gather，选取了对应label类别的概率来计算loss。

class FocalLoss(nn.Module):
    def __init__(self, alpha=0.25, gamma=2, num_classes=2, reduction='mean'):
        """
        focal_loss , -αt(1-pt)**gamma * log(pt)
        :param alpha:   balance class. list(alpha = alpha) or constant([alpha, 1-alpha, 1-alpha, ...]), default 0.25
        :param gamma:   gamma, default 2
        :param num_classes:  num classes
        :param reduction: mean or sum, default mean
        """

        super().__init__()
        self.reduction = reduction
        if isinstance(alpha, list):
            assert len(alpha) == num_classes
            self.alpha = torch.Tensor(alpha)
        else:
            assert alpha < 1   # background decay
            self.alpha = torch.zeros(num_classes)
            self.alpha[0] += alpha
            self.alpha[1:] += (1 - alpha) #  [ α, 1-α, 1-α, 1-α, 1-α, ...] size:[num_classes]
        self.gamma = gamma

    def forward(self, preds, labels):
        """
        focal_loss forward
        :param preds:   size:[N, T, C] or [T, C]    N: batch size T: video length C: num classes
        :param labels:  size:[N, T] or [T]
        :return:
        """
        preds = preds.view(-1, preds.size(-1))
        self.alpha = self.alpha.to(preds.device)
        preds_softmax = F.softmax(preds, dim=1)
        preds_logsoft = torch.log(preds_softmax)
        preds_softmax = preds_softmax.gather(1, labels.view(-1, 1))
        preds_logsoft = preds_logsoft.gather(1, labels.view(-1, 1))
        self.alpha = self.alpha.gather(0, labels.view(-1))
        loss = -torch.mul(torch.pow((1 - preds_softmax), self.gamma), preds_logsoft)
        loss = torch.mul(self.alpha, loss.t())
        if self.reduction == 'mean':
            loss = loss.mean()
        elif self.reduction == 'sum':
            loss = loss.sum()
        return loss

3.2 为什么gather后，计算bp没有只根据对应的类别的的loss计算梯度？

https://github.com/pytorch/pytorch/issues/46225

类似的

import torch
from torch.autograd import Variable

x = Variable(torch.FloatTensor([1.,1]), requires_grad=True)
div = Variable(torch.FloatTensor([0.,1]))
y = x/div

loss = y.gather(0, 1)
loss.backward()
print(x.grad)

四、解决办法

避免loss计算出现inf这种结果，如将preds_logsoft = torch.log(preds_softmax)替换为preds_logsoft = torch.log(preds_softmax+1e-8)