批量归一化
论文地址:https://arxiv.org/abs/1502.03167
批量归一化基本上是现在模型的标配了.
说实在的,到今天我也没搞明白batch normalize能够使得模型训练更稳定的底层原因,要彻底搞清楚,涉及到很多凸优化的理论,需要非常扎实的数学基础才行.
目前为止,我理解的批量归一化即把每一层输入的特征,统一变换到统一的尺度上来,避免各个特征的单位不统一的情况.即把每一个特征的分布都转变为均值为0,方差为１的分布.
然后在变换后的数据的基础上加一个线性变换.
关于batch normalize的常见问题,参考:https://zhuanlan.zhihu.com/p/55852062
对全连接层做批量归一化
我们先考虑如何对全连接层做批量归一化。通常，我们将批量归一化层置于全连接层中的仿射变换和激活函数之间。设全连接层的输入为\(\boldsymbol{u}\)，权重参数和偏差参数分别为\(\boldsymbol{W}\)和\(\boldsymbol{b}\)，激活函数为\(\phi\)。设批量归一化的运算符为\(\text{BN}\)。那么，使用批量归一化的全连接层的输出为
\]
其中批量归一化输入\(\boldsymbol{x}\)由仿射变换
\]
得到。考虑一个由\(m\)个样本组成的小批量，仿射变换的输出为一个新的小批量\(\mathcal{B} = \{\boldsymbol{x}^{(1)}, \ldots, \boldsymbol{x}^{(m)} \}\)。它们正是批量归一化层的输入。对于小批量\(\mathcal{B}\)中任意样本\(\boldsymbol{x}^{(i)} \in \mathbb{R}^d, 1 \leq i \leq m\)，批量归一化层的输出同样是\(d\)维向量
\]
并由以下几步求得。首先，对小批量\(\mathcal{B}\)求均值和方差：
\]
\]
其中的平方计算是按元素求平方。接下来，使用按元素开方和按元素除法对\(\boldsymbol{x}^{(i)}\)标准化：
\]
这里\(\epsilon > 0\)是一个很小的常数，保证分母大于0。在上面标准化的基础上，批量归一化层引入了两个可以学习的模型参数，拉伸（scale）参数 \(\boldsymbol{\gamma}\) 和偏移（shift）参数 \(\boldsymbol{\beta}\)。这两个参数和\(\boldsymbol{x}^{(i)}\)形状相同，皆为\(d\)维向量。这就是文章开头说的对特征做normalization后,再做一次线性变换
它们与\(\boldsymbol{x}^{(i)}\)分别做按元素乘法（符号\(\odot\)）和加法计算：
\]
至此，我们得到了\(\boldsymbol{x}^{(i)}\)的批量归一化的输出\(\boldsymbol{y}^{(i)}\)。
值得注意的是，可学习的拉伸和偏移参数保留了不对\(\hat{\boldsymbol{x}}^{(i)}\)做批量归一化的可能：此时只需学出\(\boldsymbol{\gamma} = \sqrt{\boldsymbol{\sigma}_\mathcal{B}^2 + \epsilon}\)和\(\boldsymbol{\beta} = \boldsymbol{\mu}_\mathcal{B}\)。我们可以对此这样理解：如果批量归一化无益，理论上，学出的模型可以不使用批量归一化。
对卷积层做批量归一化
对卷积层来说，批量归一化发生在卷积计算之后、应用激活函数之前。如果卷积计算输出多个通道，我们需要对这些通道的输出分别做批量归一化，且每个通道都拥有独立的拉伸和偏移参数，并均为标量。设小批量中有\(m\)个样本。在单个通道上，假设卷积计算输出的高和宽分别为\(p\)和\(q\)。我们需要对该通道中\(m \times p \times q\)个元素同时做批量归一化。对这些元素做标准化计算时，我们使用相同的均值和方差，即该通道中\(m \times p \times q\)个元素的均值和方差。
用个更具体点的例子总结一下就是:
对于全连接层,假设输出shape为[batch,256],那归一化即对256列的每一列求平均.
对于卷积层,假设输出shape为[batch,96,5,5],即对每个样本来说,有96个5x5的feature map,归一化在96个channel上分别做归一化,均值为batchx5x5个数的均值.
预测时的批量归一化
这时候,还有一个问题,就是模型训练好了,传入输入,计算前向传播的结果,也是要做归一化的处理的.那这时候我用的均值和方差应该是多少呢? 很显然,不应该是某个batch的样本的均值和方差,而应该是所有样本的均值和方差.因为gamma和beta的更新是不断累积的结果,而不是仅仅参考某一个batch的输入.(注意,这里的样本不是指模型的输入图片矩阵,而是指归一化层的输入,这个输入随着训练的进行是在不断变化的,而且不同的归一化层的输入是不一样的).所以,在做batch normalize的时候,我们还要维护一个值,用于估计全部样本的均值,方差.一种常见的方法是移动平均法.
可以通过下面的测试代码看一下moving_mean是如何逼近3的
momentum=0.9
moving_mean = 0.0
for epoch in range(10):
for mean in [1,2,3,4,5]:
moving_mean = momentum * moving_mean + (1.0 - momentum) * mean
print(moving_mean)
至于为何不直接对均值之和求平均,我在torch论坛提问了,目前还没回复．
现在我们来总结一下batch normalize的计算过程,然后实现它.分为训练/测试两个部分.
训练:
- 求输入x的均值
- 求输入x的方差
- 将x归一化\[\hat{\boldsymbol{x}}^{(i)} \leftarrow \frac{\boldsymbol{x}^{(i)} - \boldsymbol{\mu}_\mathcal{B}}{\sqrt{\boldsymbol{\sigma}_\mathcal{B}^2 + \epsilon}},
\] - 对归一化后的x做线性变换\[{\boldsymbol{y}}^{(i)} \leftarrow \boldsymbol{\gamma} \odot \hat{\boldsymbol{x}}^{(i)} + \boldsymbol{\beta}.
\]
测试:
- 使用移动平均所得的均值和方差,计算归一化的值
- 对归一化后的值做线性变换
那么可以写出BatchNorm的定义
def batch_norm(is_training,X,eps,gamma,beta,running_mean,running_var,alpha):
assert len(X.shape) in (2,4)
if is_training:
#X [batch,n]
if len(X.shape) == 2:
mean = X.mean(dim=0)
var = ((X-mean) ** 2).mean(dim=0)
else:
#X [batch,c,h,w]
mean = X.mean(dim=0,keepdim=True).mean(dim=2, keepdim=True).mean(dim=3, keepdim=True)
var = ((X-mean) ** 2).mean(dim=0,keepdim=True).mean(dim=2, keepdim=True).mean(dim=3, keepdim=True)
X_hat = (X - mean) / torch.sqrt(var + eps)
running_mean = alpha * mean + (1 - alpha) * running_mean
running_var = alpha * var + (1 - alpha) * running_var
else:
X_hat = (X - running_mean) / torch.sqrt(running_var + eps)
#print(gamma.shape,X_hat.shape,beta.shape)
Y = gamma * X_hat + beta #
return Y,running_mean,running_var
class BatchNorm(nn.Module):
def __init__(self,is_conv,in_channels):
super(BatchNorm,self).__init__()
#卷积层/全连接层归一化后的线性变换参数.
if not is_conv:
# x:[batch,n]
shape = (1,in_channels)
self.gamma = nn.Parameter(torch.ones(shape)) #是可学习的参数.反向传播时需要根据梯度更新.
self.beta = nn.Parameter(torch.zeros(shape)) #是可学习的参数.反向传播时需要根据梯度更新.
self.running_mean = torch.zeros(shape) #不需要求梯度.在forward时候更新.
self.running_var = torch.zeros(shape) #不需要求梯度.在forward时候更新.
else:
# x:[btach,c,h,w]
shape = (1,in_channels,1,1)
self.gamma = nn.Parameter(torch.ones(shape))
self.beta = nn.Parameter(torch.ones(shape))
self.running_mean = torch.zeros(shape)
self.running_var = torch.zeros(shape)
self.eps = 1e-5
self.momentum=0.9
def forward(self,x):
# 如果X不在内存上，将moving_mean和moving_var复制到X所在显存上
if self.running_mean.device != x.device:
self.running_mean = self.running_mean.to(x.device)
self.running_var = self.running_var.to(x.device)
# self.training继承自nn.Module,默认true,调用.eval()会设置成false
if self.training:
Y,self.running_mean,self.running_var = batch_norm(True,x,self.eps,self.gamma,self.beta,self.running_mean,self.running_var,self.momentum)
else:
Y,self.running_mean,self.running_var = batch_norm(False,x,self.eps,self.gamma,self.beta,self.running_mean,self.running_var,self.momentum)
return Y
BatchNorm继承自nn.Module,含有可学习参数gamma,beta,反向传播时会更新他们. 参数running_mean,running_var在前向传播时计算. batch_norm需要区分是卷积后的归一化还是全连接后的归一化.卷积的归一化是对每个channel单独求均值.
数据加载
batch_size,num_workers=16,2
train_iter,test_iter = learntorch_utils.load_data(batch_size,num_workers,None)
模型定义
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.conv = nn.Sequential(
nn.Conv2d(1, 6, 5), # in_channels, out_channels, kernel_size
BatchNorm(is_conv=True,in_channels=6),
nn.Sigmoid(),
nn.MaxPool2d(2, 2), # kernel_size, stride
nn.Conv2d(6, 16, 5),
BatchNorm(is_conv=True,in_channels=16),
nn.Sigmoid(),
nn.MaxPool2d(2, 2)
)
self.fc = nn.Sequential(
nn.Linear(16*4*4, 120),
BatchNorm(is_conv=False,in_channels=120),
nn.Sigmoid(),
nn.Linear(120, 84),
BatchNorm(is_conv=False,in_channels = 84),
nn.Sigmoid(),
nn.Linear(84, 10)
)
def forward(self, img):
feature = self.conv(img)
output = self.fc(feature.view(img.shape[0], -1))
return output
net = LeNet().cuda()
损失函数定义
l = nn.CrossEntropyLoss()
优化器定义
opt = torch.optim.Adam(net.parameters(),lr=0.01)
评估函数定义
def test():
acc_sum = 0
batch = 0
for X,y in test_iter:
X,y = X.cuda(),y.cuda()
y_hat = net(X)
acc_sum += (y_hat.argmax(dim=1) == y).float().sum().item()
batch += 1
print('acc:%f' % (acc_sum/(batch*batch_size)))
训练
num_epochs=5
def train():
for epoch in range(num_epochs):
train_l_sum,batch=0,0
start = time.time()
for X,y in train_iter:
X,y = X.cuda(),y.cuda() #把tensor放到显存
y_hat = net(X) #前向传播
loss = l(y_hat,y) #计算loss,nn.CrossEntropyLoss中会有softmax的操作
opt.zero_grad()#梯度清空
loss.backward()#反向传播,求出梯度
opt.step()#根据梯度,更新参数
train_l_sum += loss.item()
batch += 1
end = time.time()
time_per_epoch = end - start
print('epoch %d,train_loss %f,time %f' % (epoch + 1,train_l_sum/(batch*batch_size),time_per_epoch))
test()
train()
加了BN层以后,显存直接不够用了.但是用torch自己的nn.BatchNorm2d和nn.BatchNorm1d就没有问题.应该还是自己的对BatchNorm的实现哪里不够好.
使用torch自己的BatchＮorm的实现定义的模型如下:
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.conv = nn.Sequential(
nn.Conv2d(1, 6, 5), # in_channels, out_channels, kernel_size
nn.BatchNorm2d(6),
nn.Sigmoid(),
nn.MaxPool2d(2, 2), # kernel_size, stride
nn.Conv2d(6, 16, 5),
nn.BatchNorm2d(16),
nn.Sigmoid(),
nn.MaxPool2d(2, 2)
)
self.fc = nn.Sequential(
nn.Linear(16*4*4, 120),
nn.BatchNorm1d(120),
nn.Sigmoid(),
nn.Linear(120, 84),
nn.BatchNorm1d(84),
nn.Sigmoid(),
nn.Linear(84, 10)
)
def forward(self, img):
feature = self.conv(img)
output = self.fc(feature.view(img.shape[0], -1))
return output
net = LeNet().cuda()
训练输出如下:
epoch 1,batch_size 4,train_loss 0.194394,time 50.538379
acc:0.789400
epoch 2,batch_size 4,train_loss 0.146268,time 52.352518
acc:0.789500
epoch 3,batch_size 4,train_loss 0.132021,time 52.240710
acc:0.820600
epoch 4,batch_size 4,train_loss 0.126241,time 53.277958
acc:0.824400
epoch 5,batch_size 4,train_loss 0.120607,time 52.067259
acc:0.831800