解决pytorch半精度amp训练nan问题

Why?

如果要解决问题，首先就要明确原因：为什么全精度训练时不会nan，但是半精度就开始nan？这其实分了三种情况：

1. 计算loss 时，出现了除以0的情况
2. loss过大，被半精度判断为inf
3. 网络参数中有nan，那么运算结果也会输出nan

1&2我想放到后面讨论，因为其实大部分报nan都是第三种情况。这里来先看看3。什么情况下会出现情况3？这个讨论给出了不错的解释：

【Nan Loss with torch.cuda.amp and CrossEntropyLoss】https://link.zhihu.com/?target=https%3A//discuss.pytorch.org/t/nan-loss-with-torch-cuda-amp-and-crossentropyloss/108554/17

给大家翻译翻译：在使用ce loss 或者 bceloss的时候，会有log的操作，在半精度情况下，一些非常小的数值会被直接舍入到0，log(0)等于啥？——等于nan啊！

于是逻辑就理通了：回传的梯度因为log而变为nan->网络参数nan-> 每轮输出都变成nan。(;´Д`)

How?

问题定义清楚，那解决方案就非常简单了，只需要在涉及到log计算时，把输入从half精度转回float32：

x = x.float()
x_sigmoid = torch.sigmoid(x)

一些思考&废话

这里我接着讨论下我第一次看到nan之后，企图直接copy别人的解决方案，但解决不掉时踩过的坑。比如：

1. 修改优化器的eps

有些blog会建议你从默认的1e-8 改为 1e-3，比如这篇：【pytorch1.1 半精度训练 Adam RMSprop 优化器 Nan 问题】https://link.zhihu.com/?target=https%3A//blog.csdn.net/gwb281386172/article/details/104705195

经过上面的分析，我们就能知道为什么这种方法不行——这个方案是针对优化器的数值稳定性做的修改，而loss计算这一步在优化器之前，如果loss直接nan，优化器的eps是救不回来的（托腮）。

那么这个方案在哪些场景下有效？——在loss输出不是nan时（感觉说了一句废话）。optimizer的eps是保证在进行除法backwards时，分母不出现0时需要加上的微小量。在半精度情况下，分母加上1e-8就仿佛听君一席话，因此，需要把eps调大一点。

GradScaler是autocast的好伙伴，在官方教程上就和autocast配套使用：

from torch.cuda.amp import autocast, GradScaler
...
scaler = GradScaler()
for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()

        with autocast():
            output = model(input)
            loss = loss_fn(output, target)
        scaler.scale(loss).backward()

        scaler.step(optimizer)
        scaler.update()

具体原理不是我这篇文章讨论的范围，网上很多教程都说得很清楚了，比如这个就不错：

【Gemfield：PyTorch的自动混合精度（AMP）】https://zhuanlan.zhihu.com/p/165152789

但是我这里想讨论另一点：scaler.step(optimizer)的运行原理。

在初始化GradScaler的时候，有一个参数enabled，值默认为True。如果为True，那么在调用scaler方法时会做梯度缩放来调整loss，以防半精度状况下，梯度值过大或者过小从而被nan或者inf。而且，它还会判断本轮loss是否是nan，如果是，那么本轮计算的梯度不会回传，同时，当前的scale系数乘上backoff_factor，缩减scale的大小_。_

那么，为什么这一步已经判断了loss是不是nan，还是会出现网络损失持续nan的情况呢？

这时我们就得再往前思考一步了：为什么loss会变成nan？回到文章一开始说的：

（1）计算loss 时，出现了除以0的情况；

（2）loss过大，被半精度判断为inf；

（3）网络直接输出了nan。

(1)&(2)，其实是可以通过scaler.step(optimizer）解决的，分别由optimizer和scaler帮我们捕捉到了nan的异常。但(3)不行，(3)意味着部分甚至全部的网络参数已经变成nan了。这可能是在更之前的梯度回传过程中除以0导致的——首先【回传的梯度不是nan】，所以scaler不会捕捉异常；其次，由于使用了半精度，optimizer接收到了【已经因为精度损失而变为nan的loss】，nan不管加上多大的eps，都还是nan，所以optimizer也无法处理异常，最终导致网络参数nan。

所以3，只能通过本文一开始提出的方案来解决。其实，大部分分类问题在使用半精度时出现nan的情况都是第3种情况，也只能通过把精度转回为float32，或者在计算log时加上微小量来避免（但这样会损失精度）。

参考

【Nan Loss with torch.cuda.amp and CrossEntropyLoss】https://discuss.pytorch.org/t/nan-loss-with-torch-cuda-amp-and-crossentropyloss/108554/17

如果觉得有用，就请分享到朋友圈吧！