BN
总结
- 批量归一化固定小批量中的均值和方差，然后学习出适合的偏移和缩放。
- 可以加速收敛速度，但一般不改变模型精度。
BN代码手动实现
- 导入相关库
import torch
from torch import nn
from d2l import torch as d2l
- 定义BN层
def batch_norm(X, gamma, beta, moving_mean, moving_var, eps, momentum):
"""
:param X: 输入数据
:param gamma: γ
:param beta: β
:param moving_mean: 全局均值
:param moving_var: 全局方差
:param eps: ε
:param momentum: 冲量：用来更新或固定常量
:return:输出数据, 全局均值, 全局方差
"""
if not torch.is_grad_enabled():
# 如果是在预测模式下，直接使用传入的移动平均所得的均值和方差
X_hat = (X - moving_mean) / torch.sqrt(moving_var + eps)
else:
assert len(X.shape) in (2, 4)
if len(X.shape) == 2:
# 使用全连接层的情况：计算特征维度上的均值和方差
mean = X.mean(dim=0)
var = ((X-mean)**2).mean(dim=0)
else:
# 使用二维卷积层情况：计算通道维上（axis=1）的均值和方差
# 这里需要保持X的形状以便后面可以做广播运算
mean = X.mean(dim = (0, 2, 3), keepdim=True)
var = ((X - mean)**2).mean(dim=(0, 2, 3), keepdim=True)
# 训练模式下，用当前的均值和方差做标准化
X_hat = (X - mean) / torch.sqrt(var + eps)
# 更新移动平均的均值和方差
moving_mean = momentum * moving_mean + (1.0 - momentum) * mean
moving_var = momentum * moving_var + (1.0 - momentum) * var
Y = gamma * X_hat + beta # 缩放和移位
return Y, moving_mean.data, moving_var.data
class BatchNorm(nn.Module):
def __init__(self, num_features, num_dims):
"""
:param num_features: 全连接层的输出数量或卷积层的输出通道数
:param num_dims: 2表示全连接层，4表示卷积层
"""
super().__init__()
if num_dims == 2:
shape = (1, num_features)
else:
shape = (1, num_features, 1, 1)
# 参与求梯度和迭代的拉伸参数和偏移参数，其分别初始化为1和0
self.gamma = nn.Parameter(torch.ones(shape))
self.beta = nn.Parameter(torch.zeros(shape))
# 非模型参数的变量初始化为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:
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
- 应用BN与LeNet模型
net = nn.Sequential(
nn.Conv2d(1, 6, kernel_size=5, padding=2), BatchNorm(6, num_dims=4), nn.Sigmoid(), # [1, 6, 28, 28]
nn.AvgPool2d(2, stride=2), # [1, 6, 14, 14]
nn.Conv2d(6, 16, kernel_size=5), BatchNorm(16, num_dims=4), nn.Sigmoid(), # [1, 16, 10, 10]
nn.AvgPool2d(2, stride=2), # [1, 16, 7, 7]
nn.Flatten(), # [1, 16*5*5]
nn.Linear(16*5*5, 120), BatchNorm(120, num_dims=2), nn.Sigmoid(), # [1, 400] -->[1, 120]
nn.Linear(120, 84), BatchNorm(84, num_dims=2), nn.Sigmoid(), # [1, 120] --> [1, 84]
nn.Linear(84, 10) # [1, 82] --> [1, 10]
)
- 查看模型
X = torch.randn((1, 1, 28, 28))
for layer in net:
X = layer(X)
print(layer.__class__.__name__, 'output shape: \t', X.shape)
- 加载数据集
batch_size = 256
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
- 训练模型
import time
lr, num_epochs = 1.0, 10
start = time.perf_counter()
d2l.train_ch6(net, train_iter, test_iter, num_epochs, lr, d2l.try_gpu())
end = time.perf_counter()
print("运行耗时 %.4f s" % (end-start))
7. 查看拉伸参数gamma和偏移参数beta
net[1].gamma.reshape((-1,)), net[1].beta.reshape((-1, ))
BN代码简洁实现
- 修改模型
net = nn.Sequential(
nn.Conv2d(1, 6, kernel_size=5, padding=2),
nn.BatchNorm2d(6), nn.Sigmoid(), # [1, 6, 28, 28]
nn.AvgPool2d(2, stride=2), # [1, 6, 14, 14]
nn.Conv2d(6, 16, kernel_size=5),
nn.BatchNorm2d(16), nn.Sigmoid(), # [1, 16, 10, 10]
nn.AvgPool2d(2, stride=2), # [1, 16, 5, 5]
nn.Flatten(), # [1, 16*5*5]
nn.Linear(16 * 5 * 5, 120),
nn.BatchNorm1d(120), nn.Sigmoid(), # [1, 120]
nn.Linear(120, 84),
nn.BatchNorm1d(84), nn.Sigmoid(), # [1, 84]
nn.Linear(84, 10)) # [1, 10]
- 训练模型
d2l.train_ch6(net, train_iter, test_iter, num_epochs, lr, d2l.try_gpu())
