记录一次 NaN in Loss 的解决过程

前言

训练的模型经常数十个epoch后Loss中出现NaN，困扰了很久终于解决了，记录一下。

检测

我通常会在计算loss.backward()后optimizer.step()前，添加条件判断是否出现NaN：

if loss != loss:
    raise Exception('NaN in loss, crack!')

定位

输出参数和梯度

同样放在loss.backward()后optimizer.step()前。
用于打印网络所有层参数和梯度的最大和最小值，仅在Debug时使用：

# print grad check
v_n = []
v_v = []
v_g = []
for name, parameter in net.named_parameters():
    v_n.append(name)
    v_v.append(parameter.detach().cpu().numpy() if parameter is not None else [0])
    v_g.append(parameter.grad.detach().cpu().numpy() if parameter.grad is not None else [0])
for i in range(len(v_n)):
    if np.max(v_v[i]).item() - np.min(v_v[i]).item() < 1e-6:
        color = bcolors.FAIL + '*'
    else:
        color = bcolors.OKGREEN + ' '
    print('%svalue %s: %.3e ~ %.3e' % (color, v_n[i], np.min(v_v[i]).item(), np.max(v_v[i]).item()))
    print('%sgrad  %s: %.3e ~ %.3e' % (color, v_n[i], np.min(v_g[i]).item(), np.max(v_g[i]).item()))

部分输出如下：

 value encoder_blocks.0.KPConv.weights: -1.890e-01 ~ 9.013e-02
 grad  encoder_blocks.0.KPConv.weights: nan ~ nan
 value encoder_blocks.0.KPConv.kernel_points: -3.859e-01 ~ 4.054e-01
 grad  encoder_blocks.0.KPConv.kernel_points: 0.000e+00 ~ 0.000e+00
 value encoder_blocks.0.batch_norm.batch_norm.weight: -6.673e-03 ~ 1.823e-01
 grad  encoder_blocks.0.batch_norm.batch_norm.weight: nan ~ nan
 value encoder_blocks.0.batch_norm.batch_norm.bias: -1.067e-01 ~ 9.120e-02
 grad  encoder_blocks.0.batch_norm.batch_norm.bias: nan ~ nan

分析

首先查看是Loss中先出现NaN还是梯度中先出现NaN。
若Loss中先出现，则有可能是正向传播过程出现问题或计算Loss过程出现问题。
但我发现我的模型中梯度中首先出现NaN，说明是反向传播过程出现问题，从后往前查看梯度，找到倒数第一个NaN的地方

 value encoder.ffn0.l1.weight: -1.010e-06 ~ 1.011e-06
 grad  encoder.ffn0.l1.weight: nan ~ nan
 value encoder.ffn0.l1.bias: -9.871e-07 ~ 9.912e-07
 grad  encoder.ffn0.l1.bias: nan ~ nan
 value encoder.ffn0.l2.weight: -9.877e-07 ~ 1.031e-06
 grad  encoder.ffn0.l2.weight: nan ~ nan
 value encoder.ffn0.l2.bias: -9.813e-07 ~ 4.404e-06
 grad  encoder.ffn0.l2.bias: nan ~ nan
*value encoder.norm.weight: 3.158e-05 ~ 3.163e-05
*grad  encoder.norm.weight: -5.224e-15 ~ 9.597e-15
*value encoder.norm.bias: -2.440e-09 ~ 2.480e-09
*grad  encoder.norm.bias: -9.202e-15 ~ 1.174e-14
 value encoder.ffn1.net.0.weight: -1.017e-06 ~ 1.047e-06
 grad  encoder.ffn1.net.0.weight: -4.821e-14 ~ 6.914e-14
 value encoder.ffn1.net.0.bias: -1.174e-06 ~ 1.422e-06
 grad  encoder.ffn1.net.0.bias: -4.267e-10 ~ 6.084e-10

因此我们定位到问题是出现在 encoder.ffn0.l2 之后 encoder.norm.weight 之前，查看代码查找问题。

精确定位

知道了问题存在的位置我还是找不到问题，那就是定位得不够精确，其实可以直接使用pytorch自带的检查功能（严重影响性能，仅在Debug中使用）：

with autograd.detect_anomaly():
	# 正向传播
	# 计算Loss
	# 反向传播

当模型中出现NaN时，会抛出异常，并在打印常规的Debug信息前，提供更多的Debug信息。额外提供的信息如下：

  File "~/work/trainer.py", line 252, in train                            
    outputs = net(batch, config)                                                                      
  File "~/anaconda3/envs/pytorch1.8/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl                                                                               
    result = self.forward(*input, **kwargs)                                                           
  File "~/anaconda3/envs/pytorch1.8/lib/python3.7/site-packages/torch/nn/parallel/distributed.py", line 705, in forward                                                                            
    output = self.module(*inputs[0], **kwargs[0])                                                     
  File "~/anaconda3/envs/pytorch1.8/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl                                                                               
    result = self.forward(*input, **kwargs)                                                             
  File "~/anaconda3/envs/pytorch1.8/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl                                                                               
    result = self.forward(*input, **kwargs)                                                           
  File "~/work/models/crossvitmodule.py", line 672, in forward                  
    x = list(x)                                                                                       
  File "~/work/models/crossvitmodule.py", line 663, in <lambda>                 
    x = map(lambda u, v: u + self.pos_embedding(u, v), x, point)                                      
  File "~/anaconda3/envs/pytorch1.8/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl                                                                               
    result = self.forward(*input, **kwargs)                                                           
  File "~/work/models/crossvitmodule.py", line 423, in forward                  
    point = self.proj2(point)                                                                         
  File "~/anaconda3/envs/pytorch1.8/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl                                                                               
    result = self.forward(*input, **kwargs)                                                           
  File "~/work/models/crossvitmodule.py", line 404, in forward                  
    return self.leaky_relu(self.norm(self.linear(x)))                                                 
  File "~/anaconda3/envs/pytorch1.8/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl                                                                               
    result = self.forward(*input, **kwargs)                                                           
  File "~/work/models/crossvitmodule.py", line 29, in forward                   
    std = torch.var(x, dim=1, unbiased=False, keepdim=True).sqrt()                        
 (function _print_stack)   

直接定位到std = torch.var(x, dim=1, unbiased=False, keepdim=True).sqrt()这一行。
原代码为手动实现的LayerNorm：

class LayerNorm(nn.Module):
    def __init__(self, dim, eps = 1e-5):
        super().__init__()
        self.eps = eps
        self.g = nn.Parameter(torch.ones(1, dim, 1, 1))
        self.b = nn.Parameter(torch.zeros(1, dim, 1, 1))

    def forward(self, x):
        std = torch.var(x, dim = 1, unbiased = False, keepdim = True).sqrt()
        mean = torch.mean(x, dim = 1, keepdim = True)
        return (x - mean) / (std + self.eps) * self.g + self.b

当x为0时在计算标准差时正向传播需要计算sqrt(x)
而反向传播时需要计算sqrt(x)的微分1/(2*sqrt(x))此时需要确保x != 0

解决方案

防止除0即可：

std = torch.sqrt(torch.var(x, dim = 1, unbiased = False, keepdim = True) + self.eps)