1. 从零实现
下面,我们从头开始实现一个具有张量的批量规范化层。
import torch
from torch import nn
from d2l import torch as d2l
# X是输入,gamma, beta是两个可以学习的参数,moving_mean, moving_var是整个数据集的均值和方差,
# 而不是小批量的均值和方差,做inference(test)的时候用,eps防止除零错误,通常不改变值
# momentum是来更新moving_mean, moving_var的
def batch_norm(X, gamma, beta, moving_mean, moving_var, eps, momentum):
# 通过is_grad_enabled来判断当前模式是训练模式还是预测模式
# 训练和预测时 BN的过程不一样
if not torch.is_grad_enabled(): # 不用计算梯度,就代表不是在做训练,是做inference
# 推理就是forward,不需要反向传播,所以不需要计算梯度
# 如果是在预测模式下,直接使用传入的移动平均所得的均值和方差
X_hat = (X - moving_mean) / torch.sqrt(moving_var + eps) # 推理时用的全局的mean和var
# 为什么不能用一个批量的均值和方差?
# 答:因为做inference的时候,可能没有一个批量,可能就一个样本进来
else: # 训练模式下
assert len(X.shape) in (2, 4) # X.shape要么为2,表示全连接层,要么为4,表示卷积层
if len(X.shape) == 2:# (batch_size,feature)
# 使用全连接层的情况,计算特征维上的均值和方差
mean = X.mean(dim=0) # 表示跨行(dim=0)求平均,也就是对同一列元素求均值
var = ((X - mean) ** 2).mean(dim=0) # 因此方差和mean都是一个行向量
else: # (batch_size,channels,height,width)
# 使用二维卷积层的情况,计算通道维上(axis=1)的均值和方差。
# 这里我们需要保持X的形状以便后面可以做广播运算
mean = X.mean(dim=(0, 2, 3), keepdim=True)
# dim=(0, 2, 3)表示对每一个通道的所有批量*所有高宽的像素全部来求均值
# 换个角度想:dim=1的话表示哭啊通道求平均,也就是每一行求平均
# dim=(0, 2, 3)与之相反,那么就是每一列求平均
# 最后求出来的结果是 1 * n * 1 * 1,因为keepdim=True
var = ((X - mean) ** 2).mean(dim=(0, 2, 3), keepdim=True)
# X - mean用了广播机制,X为2 *2 * 3* 3,mean为 1 * 2 * 1 * 1,可以进行广播机制
# 最后求出来的结果也是 1 * n * 1 * 1
# 训练模式下,用当前的均值和方差做标准化
X_hat = (X - mean) / torch.sqrt(var + eps) # 训练时使用的是当前的小批量中的mean和方差
# 更新移动平均的均值和方差:
# 每得到一次小批量的均值和方差,就去更新全局的数据,这样会无限逼近真实的在数据集上的均值和方差
moving_mean = momentum * moving_mean + (1.0 - momentum) * mean
moving_var = momentum * moving_var + (1.0 - momentum) * var
# 返回的Y就是对X做了线性变化之后的数据
Y = gamma * X_hat + beta # 缩放和移位
return Y, moving_mean.data, moving_var.data
我们现在可以创建一个正确的BatchNorm层。 这个层将保持适当的参数:拉伸gamma和偏移beta,这两个参数将在训练过程中更新。 此外,我们的层将保存均值和方差的移动平均值,以便在模型预测期间随后使用。
撇开算法细节,注意我们实现层的基础设计模式。
通常情况下,我们用一个单独的函数定义其数学原理,比如说batch_norm。 然后,我们将此功能集成到一个自定义层中,其代码主要处理数据移动到训练设备(如GPU)、分配和初始化任何必需的变量、跟踪移动平均线(此处为均值和方差)等问题。
为了方便起见,我们并不担心在这里自动推断输入形状,因此我们需要指定整个特征的数量。 不用担心,深度学习框架中的批量规范化API将为我们解决上述问题,我们稍后将展示这一点。
class BatchNorm(nn.Module):
# num_features:完全连接层的输出数量或卷积层的输出通道数。
# num_dims:2表示完全连接层,4表示卷积层
def __init__(self, num_features, num_dims):
super().__init__()
# 为什么要定义shape的原因?
# 答:gamma,beta,moving_mean和moving_var的shape都是和shape一样
if num_dims == 2: # 完全连接层
shape = (1, num_features)
else: # 卷积层
shape = (1, num_features, 1, 1)
# 参与求梯度和迭代的拉伸和偏移参数,分别初始化成1和0
# gamma如果初始化为0,那么gamma乘以X_hat等于0,训练不动
self.gamma = nn.Parameter(torch.ones(shape))
self.beta = nn.Parameter(torch.zeros(shape))
# gamma和 beta是需要被迭代的,moving_mean和moving_var则不需要被迭代
# 所以gamma和 beta放在nn.Parameter中,moving_mean和moving_var就不需要放在里面
# 放入nn.Parameter中的变量,会随着网络的优化器一起更新迭代
# 非模型参数的变量初始化为0和1
self.moving_mean = torch.zeros(shape)
self.moving_var = torch.ones(shape)
def forward(self, X):
# 如果X不在内存上,将moving_mean和moving_var
# 复制到X所在显存上
if self.moving_mean.device != X.device:
# 因为moving_mean和moving_var 没有放在nn.Parameter中,也就是没放入网络net中,
# 而网络是在GPU上,所以需要移动这两个参数
self.moving_mean = self.moving_mean.to(X.device)
self.moving_var = self.moving_var.to(X.device)
# 保存更新过的moving_mean和moving_var
Y, self.moving_mean, self.moving_var = batch_norm(
X, self.gamma, self.beta, self.moving_mean,
self.moving_var, eps=1e-5, momentum=0.9)
return Y
2. 使用批量规范化层的 LeNet
为了更好理解如何应用BatchNorm,下面我们将其应用(于LeNet模型)。 回想一下,批量规范化是在卷积层或全连接层之后、相应的激活函数之前应用的。
net = nn.Sequential(
nn.Conv2d(1, 6, kernel_size=5), BatchNorm(6, num_dims=4), nn.Sigmoid(),
nn.AvgPool2d(kernel_size=2, stride=2),
nn.Conv2d(6, 16, kernel_size=5), BatchNorm(16, num_dims=4), nn.Sigmoid(),
nn.AvgPool2d(kernel_size=2, stride=2), nn.Flatten(),
nn.Linear(16*4*4, 120), BatchNorm(120, num_dims=2), nn.Sigmoid(),
nn.Linear(120, 84), BatchNorm(84, num_dims=2), nn.Sigmoid(),
nn.Linear(84, 10))
和以前一样,我们将在Fashion-MNIST数据集上训练网络。 这个代码与我们第一次训练LeNet时几乎完全相同,主要区别在于学习率大得多。
lr, num_epochs, batch_size = 1.0, 10, 256
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
d2l.train_ch6(net, train_iter, test_iter, num_epochs, lr, d2l.try_gpu())
运行结果:
让我们来看看从第一个批量规范化层中学到的拉伸参数gamma和偏移参数beta。
net[1].gamma.reshape((-1,)), net[1].beta.reshape((-1,))
3. 简洁实现
除了使用我们刚刚定义的BatchNorm,我们也可以直接使用深度学习框架中定义的BatchNorm。 该代码看起来几乎与我们上面的代码相同。
net = nn.Sequential(
# nn.BatchNorm2d(6)传入的参数是:完全连接层的输出数量或卷积层的输出通道数
nn.Conv2d(1, 6, kernel_size=5), nn.BatchNorm2d(6), nn.Sigmoid(),
nn.AvgPool2d(kernel_size=2, stride=2),
nn.Conv2d(6, 16, kernel_size=5), nn.BatchNorm2d(16), nn.Sigmoid(),
nn.AvgPool2d(kernel_size=2, stride=2), nn.Flatten(),
nn.Linear(256, 120), nn.BatchNorm1d(120), nn.Sigmoid(),
nn.Linear(120, 84), nn.BatchNorm1d(84), nn.Sigmoid(),
nn.Linear(84, 10))
下面,我们使用相同超参数来训练模型。 请注意,通常高级API变体运行速度快得多,因为它的代码已编译为C++或CUDA,而我们的自定义代码由Python实现。
d2l.train_ch6(net, train_iter, test_iter, num_epochs, lr, d2l.try_gpu())
4. 一些Q&A
Q1: 之前讲xavier的时候,也讲过类似的normalization的东西,这个和这里的BN有什么区别?
A1:讲xavier是选取好的初始化,对权重进行好的初始化(对权重分布的均值和方差进行控制),使模型训练比较稳定,但不能保证之后。
而BN则是保证在整个模型训练的时候对feature map进行均值和方差的控制,即:在每一个全连接层和卷积层前面或后面做normalization。
Q2:批量归一化是不是和权重衰退的作用类似?
A2:不类似。权重衰退是指在更新参数(权重)的时候,把权重除了一个值,使其变小,但批量归一化是不会对前面一层的权重有影响,他只是对上一层输出的数据进行规范化(均值和方差的控制)。
Q3: batch norm 能用在mlp中吗?
A3:可以的,但是batch norm对比较深的神经网络相对有用,因为只有当网络比较深的时候,才会出现上面更新快,而靠近输入数据的更新慢的情况,此时用BN效果会好。浅层MLP加上BN效果不见得很好。
Q4: BN是做了线性变换,和加一个线性层有什么区别?
A4:如果我就做一般的线性变换,那我对这个数据做怎样的处理呢?我也不知道要这个线性层做什么才能使得数值稳定,我加一个线性层能学到的参数不一定能学得这么好。而BN就相当于一个先验知识,告诉你这么做是可以使得训练更稳定的。并且,一般的线性变化Linear是同一个样本的各个维度,而BN是针对不同样本的同一特征。
Q5:为什么加了BN之后,收敛时间变短?
A5:BN会使得梯度的值变大一些,而且每个层的梯度的值会差不多,即不同层之间的梯度稳定,这样导致可以使用更大的学习率,就会使得对权重的更新变快。
有一种解释是,BN之后地形图更加平坦,更容易收敛。
Q6:layer normalization和 batch normalization 的异同。
A6: batch normalization是在feature维度对样本做normalization(对不同样本的同一特征维度做规范化)。
layer normalization就不是像batch normalization跨样本做规范化,而是在每一个样本里面做规范化,也就是对同一个样本的所有特征进行规范化。
Q7:batch size 是把显存占满好?还是利用率gpu-util 100%就可以了?还是需要同时达到100%?
A7:增加bath size ,去看一下每一秒能处理的样本数,如果增到一个程度发现每一秒能处理的样本数不会再增加,就可以停止增加batch size。但这不针对LeNet,因为LeNet可以有很大的batch size,增大到1000都可以,但是收敛会有问题。