最近在研究yolo的算法源码,在调试过程中发现中间层的BatchNorm2d的结果竟然出现了Nan。

【深度学习】完全解读BatchNorm2d归一化算法原理_深度学习

【深度学习】完全解读BatchNorm2d归一化算法原理_人工智能_02

【深度学习】完全解读BatchNorm2d归一化算法原理_算法_03

根据模型处于训练阶段或测试阶段,参数trainning和track_running_states有4种组合方式。

  • trainning = True,track_running_states = True,模型处于训练阶段,表示每作一次归一化,模型都需要更新参数均值和方差,即更新参数 running_mean 和 running_var 。
  • trainning = True,track_running_stats = False,模型处于训练阶段,表示对新的训练数据进行归一化时,不更新模型的均值和方差,这种设置是错误的,因为不能很好的描述全局的数据统计特性。
  • trainning = False,track_running_stats = True,模型处于测试阶段,表示模型在归一化测试数据时,需要考虑模型的均值和方差,但是不更新模型的均值和方差。
  • trainning = False,track_running_stats = False,模型处于测试阶段,表示模型在归一化测试数据时,不考虑模型的均值和方差,这种设置是错误的,归一化的结果会造成统计特性的偏移。

由上面4种组合参数的介绍,正确的参数设置应为:

训练阶段:trainning = True,track_running_stats = True

测试阶段:training = False,track_running_stats = True

脑壳是不是一头雾水,下面用例子说明参数的含义:

1. 训练阶段的归一化实例

初始化训练阶段的归一化模型:

m3 = nn.BatchNorm2d(3, eps=0, momentum=0.5, affine=True, track_running_stats=True).cuda()
# 为了方便验证,设置模型参数的值
m3.running_mean = (torch.ones([3])*4).cuda()     # 设置模型的均值是4
m3.running_var = (torch.ones([3])*2).cuda()      # 设置模型的方差是2

# 查看模型参数的值
print('trainning:',m3.training)
print('running_mean:',m3.running_mean)
print('running_var:',m3.running_var)
# gamma对应模型的weight,默认值是1
print('weight:',m3.weight)
# gamma对应模型的bias,默认值是0
print('bias:',m3.bias)

#>
trainning: True
running_mean: tensor([4., 4., 4.], device='cuda:0')
running_var: tensor([2., 2., 2.], device='cuda:0')
weight: Parameter containing:
tensor([1., 1., 1.], device='cuda:0', requires_grad=True)
bias: Parameter containing:
tensor([0., 0., 0.], device='cuda:0', requires_grad=True)

生成批量数据为1,通道为3,均值为0方差为1的416行416列输入数据:

# 生成通道3,416行416列的输入数据
torch.manual_seed(21)
input3 = torch.randn(1, 3, 416, 416).cuda()
# 输出第一个通道的数据
input3[0][0]

#>
tensor([[-0.2386, -1.0934,  0.1558,  ..., -0.3553, -0.1205, -0.3859],
        [ 0.2582,  0.2833,  0.7942,  ...,  1.1228,  0.3332, -1.2364],
        [-0.8235, -1.1512, -0.5026,  ...,  0.9393, -0.5026, -0.4719],
        ...,
        [-0.2843, -1.3638, -0.4599,  ...,  1.6502,  0.4864, -0.1804],
        [ 0.3813, -0.6426,  0.4879,  ...,  2.7496,  1.8501,  1.7092],
        [ 0.8221, -0.5702,  0.1705,  ...,  1.0553,  1.0248,  0.5127]],
       device='cuda:0')

对上面的数据进行归一化:

# 数据归一化
output3 = m3(input3)
# 输出归一化后的第一个通道的数据
output3[0][0]

#>
tensor([[-0.2427, -1.0955,  0.1508,  ..., -0.3592, -0.1249, -0.3897],
        [ 0.2529,  0.2779,  0.7876,  ...,  1.1154,  0.3277, -1.2382],
        [-0.8262, -1.1531, -0.5061,  ...,  0.9323, -0.5061, -0.4755],
        ...,
        [-0.2884, -1.3652, -0.4635,  ...,  1.6416,  0.4805, -0.1847],
        [ 0.3757, -0.6458,  0.4820,  ...,  2.7383,  1.8410,  1.7004],
        [ 0.8154, -0.5735,  0.1654,  ...,  1.0480,  1.0176,  0.5067]],
       device='cuda:0', grad_fn=<SelectBackward>)

为了理解BatchNorm2d的函数实现,我们编写此函数的算法实现,比对归一化结果。

因为trainning = True,track_running_stats = True,我们需要更新模型的均值和方差:

【深度学习】完全解读BatchNorm2d归一化算法原理_人工智能_04

【深度学习】完全解读BatchNorm2d归一化算法原理_深度学习_05

为模型更新前的均值或方差,代码计算更新后的均值和方差:

# 计算更新后的均值和方差
momentum = m3.momentum    # 更新参数
# 更新均值
ex_new = (1 - momentum) * ex_old + momentum * obser_mean
# 更新方差
var_new = (1 - momentum) * var_old + momentum * obser_var
# 打印
print('ex_new:',ex_new)
print('var_new:',var_new)

#>
ex_new: tensor([2.0024, 2.0015, 2.0007], device='cuda:0')
var_new: tensor([1.5024, 1.4949, 1.5012], device='cuda:0')

我们不调用归一化函数,自己编写训练阶段的归一化代码:

# 输入数据的均值
obser_mean = torch.Tensor([input3[0][i].mean() for i in range(3)]).cuda()
# 输入数据的方差
obser_var = torch.Tensor([input3[0][i].var() for i in range(3)]).cuda()
# 编码归一化
output3_source = (input3[0][0] - obser_mean[0])/(pow(obser_var[0] + m3.eps,0.5))
output3_source

#>
tensor([[-0.2427, -1.0955,  0.1508,  ..., -0.3592, -0.1249, -0.3897],
        [ 0.2529,  0.2779,  0.7876,  ...,  1.1154,  0.3277, -1.2382],
        [-0.8262, -1.1531, -0.5061,  ...,  0.9323, -0.5061, -0.4755],
        ...,
        [-0.2884, -1.3652, -0.4635,  ...,  1.6416,  0.4805, -0.1847],
        [ 0.3757, -0.6458,  0.4820,  ...,  2.7383,  1.8410,  1.7004],
        [ 0.8154, -0.5735,  0.1654,  ...,  1.0480,  1.0176,  0.5067]],
       device='cuda:0')

结果一致,我们输出模型的running_mean和running_var:

m3.running_mean,m3.running_var

#>
(tensor([2.0024, 2.0015, 2.0007], device='cuda:0'),
 tensor([1.5024, 1.4949, 1.5012], device='cuda:0'))

发现模型的running_mean和running_var和我们计算的更新均值与方差结果一致,同时通过代码我们也知道模型的running_mean和running_var是在forward()操作中更新的,训练阶段的算法实现就介绍到这,下面介绍下测试阶段归一化的算法实现。

2. 测试阶段的归一化实例

初始化归一化模型,并设置模型处于测试阶段:

# 初始化模型,并设置模型处于测试阶段
import torch
import torch.nn as nn
m3 = nn.BatchNorm2d(3, eps=0, momentum=0.5, affine=True, track_running_stats=True).cuda()
# 测试阶段
m3.eval()
# 为了方便验证,设置模型参数的值
m3.running_mean = (torch.ones([3])*4).cuda()     # 设置模型的均值是4
m3.running_var = (torch.ones([3])*2).cuda()      # 设置模型的方差是2

# 查看模型参数的值
print('trainning:',m3.training)
print('running_mean:',m3.running_mean)
print('running_var:',m3.running_var)
# gamma对应模型的weight,默认值是1
print('weight:',m3.weight)
# gamma对应模型的bias,默认值是0
print('bias:',m3.bias)

#>
trainning: False
running_mean: tensor([4., 4., 4.], device='cuda:0')
running_var: tensor([2., 2., 2.], device='cuda:0')
weight: Parameter containing:
tensor([1., 1., 1.], device='cuda:0', requires_grad=True)
bias: Parameter containing:
tensor([0., 0., 0.], device='cuda:0', requires_grad=True)

生成3通道,416行416列的输入数据

# 初始化输入数据,并计算输入数据的均值和方差
# 生成通道3,416行416列的输入数据
torch.manual_seed(21)
input3 = torch.randn(1, 3, 416, 416).cuda()
# 输入数据的均值
obser_mean = torch.Tensor([input3[0][i].mean() for i in range(3)]).cuda()
# 输入数据的方差
obser_var = torch.Tensor([input3[0][i].var() for i in range(3)]).cuda()
# 打印
print('obser_mean:',obser_mean)
print('obser_var:',obser_var)

#>
obser_mean: tensor([0.0047, 0.0029, 0.0014], device='cuda:0')
obser_var: tensor([1.0048, 0.9898, 1.0024], device='cuda:0')

归一化输入数据,并打印第一个通道的数据

# 数据归一化
output3 = m3(input3)
# 输出归一化后的第一个通道的数据
output3[0][0]

#>
tensor([[-2.9971, -3.6016, -2.7182,  ..., -3.0797, -2.9136, -3.1013],
        [-2.6459, -2.6281, -2.2668,  ..., -2.0345, -2.5928, -3.7027],
        [-3.4107, -3.6424, -3.1838,  ..., -2.1642, -3.1838, -3.1621],
        ...,
        [-3.0295, -3.7928, -3.1536,  ..., -1.6615, -2.4845, -2.9560],
        [-2.5588, -3.2828, -2.4834,  ..., -0.8842, -1.5202, -1.6199],
        [-2.2471, -3.2316, -2.7078,  ..., -2.0822, -2.1038, -2.4659]],
       device='cuda:0', grad_fn=<SelectBackward>)

自己编写测试阶段的归一化代码,结果与调用BatchNorm2d函数结果一致。

# 归一化函数实现
output3_source = (input3[0][0] - m3.running_mean[0])/(pow(m3.running_var[0] + m3.eps,0.5))
output3_source

#>
tensor([[-2.9971, -3.6016, -2.7182,  ..., -3.0797, -2.9136, -3.1013],
        [-2.6459, -2.6281, -2.2668,  ..., -2.0345, -2.5928, -3.7027],
        [-3.4107, -3.6424, -3.1838,  ..., -2.1642, -3.1838, -3.1621],
        ...,
        [-3.0295, -3.7928, -3.1536,  ..., -1.6615, -2.4845, -2.9560],
        [-2.5588, -3.2828, -2.4834,  ..., -0.8842, -1.5202, -1.6199],
        [-2.2471, -3.2316, -2.7078,  ..., -2.0822, -2.1038, -2.4659]],
       device='cuda:0')

打印模型的running_mean和running_var

# 查看模型的running_mean和running_var
print(m3.running_mean,m3.running_var)

#>
tensor([4., 4., 4.], device='cuda:0') tensor([2., 2., 2.], device='cuda:0')

由结果可知,执行测试阶段的froward函数后,模型的running_mean和running_var不改变。

小结

3. 小结

由上面例子可知:

当trainning = True,track_running_stats = True,训练阶段改变了模型的running_mean和running_var,归一化算法的均值和方差采用了模型更新前的running_mean和 running_var 。

当trainning = False,track_running_stats = True,测试阶段不改变了模型的running_mean和running_var,归一化算法的均值和方差采用了模型的running_mean和 running_var 。

其他两种情况(trainning = True,track_running_stats = False 和 trainning = False ,track_running_stats = False),小伙伴可以用上述例子去验证这两种情况的算法实现,因为这两种情况会产生较大的偏差,这里不作介绍了。

回到本文的第一个问题,为什么归一化后会出现Nan,原来由于前面的失误,造成最后一个通道的方差小于 0 ,如下红框标记的图:

【深度学习】完全解读BatchNorm2d归一化算法原理_方差_06


【深度学习】完全解读BatchNorm2d归一化算法原理_算法_07