pytorch和mmdetection训练出现nan的解决方案

本文主要是收集了一些在使用pytorch自带的amp下loss 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

给大家翻译翻译：在使用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 问题

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

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

2. 聊聊amp的GradScaler

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）

但是我这里想讨论另一点：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

1. 问题描述

在使用mmdetection训练detectors_htc_r50模型时，出现了显存不够用的问题，于是考虑采用半精度训练，但是在训练时出现了这个问题：

图1. 问题截图

我的代码环境是：

1. 系统：ubuntu20.04
2. GPU：NVIDIA RTX 2080Ti，显存：11GB
3. python：3.7
4. CUDA：11.4
5. pytorch：1.7.1
6. mmcv-full：1.1.5

在网上查了很久，最终探索出了一个非常简单的解决方法。

在你的config.py文件里，要使用fp16训练需要有这样一句话，其中loss_scale默认为512.，它是loss的缩放因子，它越大，loss就越大，因此出现了loss过大变为nan的情况。

fp16 = dict(loss_scale=512.) # 解决方案：缩小loss_scale，在我的问题中，loss_scale=64.就可以正常运行了

解决方案：缩小loss_scale，在我的问题中，loss_scale=64.就可以正常运行了。

2. 解决过程

我分析了图1所示的输出，发现rpn_head的loss并没有出现问题，出现问题的部分在roi_head的bbox_head部分。于是开始从tools/train.py向上找出现问题的具体位置。

（1）根据网上[1,2,3]的经验，我猜测可能是由于在计算loss的时候，出现了送到log()函数里的值为0的情况。于是查看了mmdet/models/losses/cross_entropy_loss.py中的CrossEntropyLoss类，输出了其中送到log_softmax函数里的cls_score项，发现输出有很多nan。这说明并不是送到log()函数里的值为0的问题，而是模型输出的预测结果就已经是nan了；

（2）然后我又顺着train.py一点一点往上查，最终在mmdet/apis/train.py中发现fp16的设置部分：

# fp16 setting
    fp16_cfg = cfg.get('fp16', None)
    if fp16_cfg is not None:
        optimizer_config = Fp16OptimizerHook(
            **cfg.optimizer_config, **fp16_cfg, distributed=distributed)
    elif distributed and 'type' not in cfg.optimizer_config:
        optimizer_config = OptimizerHook(**cfg.optimizer_config)
    else:
        optimizer_config = cfg.optimizer_config

在Fp16OptimizerHook类的定义（mmdet/core/fp16/hooks.py）中发现了一句注释：

Args:
        loss_scale (float): Scale factor multiplied with loss.

说明loss_scale和loss是相乘的关系，于是就尝试了缩小loss_scale，果然解决了问题。

关于半精度训练的更多知识，可以参考[4]。

 

参考

[1]Detectors' loss is nan with FP16 · Issue #3605 · open-mmlab/mmdetection (github.com)

[2]解决pytorch半精度amp训练nan问题 - 知乎 (zhihu.com)

[3]Nan Loss with torch.cuda.amp and CrossEntropyLoss - mixed-precision - PyTorch Forums

[4]全网最全-混合精度训练原理 - 知乎 (zhihu.com)