PyTorch踩坑指南（1）nn.BatchNorm2d()函数

大家好，又见面了，我是你们的朋友全栈君。

前言

最近在研究深度学习中图像数据处理的细节，基于的平台是PyTorch。心血来潮，总结一下，好记性不如烂笔头。

Batch Normalization

对于2015年出现的Batch Normalization¹，2018年的文章Group Normalization²在Abstract中总结得言简意赅，我直接copy过来。

Batch Normalization (BN) is a milestone technique in the development of deep learning, enabling various networks to train. However, normalizing along the batch dimension introduces problems — BN’s error increases rapidly when the batch size becomes smaller, caused by inaccurate batch statistics estimation. This limits BN’s usage for training larger models and transferring features to computer vision tasks including detection, segmentation, and video, which require small batches constrained by memory consumption.

机器学习中，进行模型训练之前，需对数据做归一化处理，使其分布一致。在深度神经网络训练过程中，通常一次训练是一个batch，而非全体数据。每个batch具有不同的分布产生了internal covarivate shift问题——在训练过程中，数据分布会发生变化，对下一层网络的学习带来困难。Batch Normalization强行将数据拉回到均值为0，方差为1的正太分布上，一方面使得数据分布一致，另一方面避免梯度消失。

结合图1，说明Batch Normalization的原理。假设在网络中间经过某些卷积操作之后的输出的feature maps的尺寸为N×C×W×H，5为batch size（N），3为channel（C），W×H为feature map的宽高，则Batch Normalization的计算过程如下。
在这里插入图片描述

图 1

1.每个batch计算同一通道的均值 $\mu$ ，如图取channel 0，即 $c = 0$ （红色表示）
$\mu = \frac{\sum\limits_{n=0}^{N-1}\sum\limits_{w=0}^{W-1} \sum\limits_{h=0}^{H-1} X[n, c, w, h]}{N×W×H}$
2.每个batch计算同一通道的方差 $σ^2$
$σ^2 = \frac{\sum\limits_{n=0}^{N-1}\sum\limits_{w=0}^{W-1} \sum\limits_{h=0}^{H-1} (X[n, c, w, h]-\mu)^2}{N×W×H}$
3.对当前channel下feature map中每个点 $x$ ，索引形式 $X [n, c, w, h]$ ，做归一化
$x^{‘}=\frac{(x-\mu)}{\sqrt{σ^2+\epsilon}}$
4.增加缩放和平移变量 $\gamma$ 和 $\beta$ （可学习的仿射变换参数），归一化后的值
$y=\gamma x^{‘}+\beta$
简化公式：
$y=\frac{x-\mu}{\sqrt{\sigma^2+\epsilon}}\gamma +\beta$
原文中的算法描述如下，

注：上图¹所示 $m$ 就是 $N * W * H$

PyTorch的nn.BatchNorm2d()函数

理解了Batch Normalization的过程，PyTorch里面的函数就参考其文档³用就好。
BatchNorm2d()内部的参数如下：

num_features：一般情况下输入的数据格式为batch_size * num_features * height * width，即为特征数，channel数
eps：分母中添加的一个值，目的是为了计算的稳定性，默认：1e-5
momentum：一个用于运行过程中均值和方差的一个估计参数，默认值为 $0.1$ ； $\hat{x}_{new} =(1−momentum) × \hat{x} +momentum×x_t$ ，其中 $\hat{x}$ 是估计值， $x_t$ 是新的观测值
affine：当设为true时，给定可以学习的系数矩阵 $\gamma$ 和 $\beta$

Show me the codes

import torch
import torch.nn as nn

def checkBN(debug = False):
    # parameters
    N = 5 # batch size
    C = 3 # channel
    W = 2 # width of feature map
    H = 2 # height of feature map
    # batch normalization layer
    BN = nn.BatchNorm2d(C,affine=True) #gamma和beta, 其维度与channel数相同
    # input and output
    featuremaps = torch.randn(N,C,W,H)
    output = BN(featuremaps)
    # checkout
    ###########################################
    if debug:
        print("input feature maps：\n",featuremaps)
        print("normalized feature maps: \n",output)
    ###########################################
    
    # manually operation, the first channel
    X = featuremaps[:,0,:,:]
    firstDimenMean = torch.Tensor.mean(X)
    firstDimenVar = torch.Tensor.var(X,False) #Bessel's Correction贝塞尔校正不被使用
    
    BN_one = ((input[0,0,0,0] - firstDimenMean)/(torch.pow(firstDimenVar+BN.eps,0.5) )) * BN.weight[0] + BN.bias[0]
    print('+++'*15,'\n','manually operation: ', BN_one)
    print('==='*15,'\n','pytorch result: ', output[0,0,0,0])
    
if __name__=="__main__":
    checkBN()

可以看出手算的结果和PyTorch的nn.BatchNorm2d的计算结果一致。

+++++++++++++++++++++++++++++++++++++++++++++
 manually operation:  tensor(-0.0327, grad_fn=<AddBackward0>)
=============================================
 pytorch result:  tensor(-0.0327, grad_fn=<SelectBackward>)

贝塞尔校正

代码中出现，求方差时是否需要贝塞尔校正，即从样本方差到总体方差的校正。
方差公式从，
$\sigma^2 = \frac{\sum\limits_{i=0}^{N-1} (x_i-mean(x))^2}{N}$
变成（基于样本的总体方差的无偏估计）,
$\sigma^2 = \frac{\sum\limits_{i=0}^{N-1} (x_i-mean(x))^2}{N-1}$