深度学习代码|Batch Normalization批归一化的代码实现

---

一、导入相关库

import torch
from torch import nn
from labml_helpers.module import Module

二、批量归一化层BatchNorm

（一）理论基础

1. 当输入$X\in {\mathbb{R}}^{B\times C\times H\times W}$是一批图像表示时：$BN(X)=\gamma \frac{X-{\mathbb{E}}_{B,H,W}[X]}{\sqrt{Va{r}_{B,H,W}[X]+ϵ}}+\beta$
2. 当输入$X\in {\mathbb{R}}^{B\times C}$是一批嵌入时：$BN(X)=\gamma \frac{X-{\mathbb{E}}_B[X]}{\sqrt{Va{r}_B[X]+ϵ}}+\beta$
3. 当输入$X\in {\mathbb{R}}^{B\times C\times L}$是一批序列嵌入时：$BN(X)=\gamma \frac{X-{\mathbb{E}}_{B,L}[X]}{\sqrt{Va{r}_{B,L}[X]+ϵ}}+\beta$

B：批次大小，C：通道数/特征数，H：高度，W：宽度，L：顺序

（二）代码实现

• __init__函数：
  • channels：输入中的特征数
  • eps：$ϵ$， 用于数值稳定性
  • momentum：更新指数移动平均值的动量参数
  • affine：是否使用可学习的拉伸和偏移参数
  • track_running_stats：计算移动平均线还是均值和方差

动量参数控制着对当前批次计算得到的均值和方差与之前计算的指数移动平均的均值和方差之间的权衡。

• forward 方法实现了批标准化的前向计算过程。
  • 将输入张量 x 重塑为 (batch_size, channels, -1) 的形状，以便在每个通道上计算均值和方差。
  • 计算当前批次的均值和方差，或者使用保存的指数移动平均的均值和方差。
  • 接着对输入进行标准化，并根据是否使用可学习参数来对标准化后的结果进行拉伸和偏移。
  • 最后将结果重塑回原来的形状并返回。

• register_buffer 是 PyTorch 中 Module 类的方法之一，用于将张量缓冲区注册到模型中，以便在模型的状态中进行跟踪。这些缓冲区在模型的训练和推理过程中保持不变，可以用于存储模型的状态信息，如均值、方差等。
• 使用 register_buffer 方法的好处是，这些缓冲区不会被包括在模型的参数列表中，因此在模型进行优化时不会更新它们的值。这使得它们适用于存储在训练过程中需要持久化的状态信息，而不会被认为是需要训练的参数。

class BatchNorm(Module):
	def __init__(self,channels:int,*,
							eps:float=1e-5,momentum:float=0.1,
							affine:bool=True,track_running_stats:bool=True):
		super().__init__()
		self.channels=channels
		self.eps=eps
		self.momentum=momentum
		self.affine=affine
		self.track_running_stats=track_running_stats
		#如果 affine=True，则创建可学习的拉伸和偏移参数 scale 和 shift
		if self.affine:
			self.scale=nn.Parameter(torch.ones(channels))
			self.shift=nn.Parameter(torch.zeros(channels))
		#如果 track_running_stats=True，则创建exp_mean 和 exp_var，用于保存指数移动平均的均值和方差
		if self.track_running_stats:
			self.register_buffer('exp_mean',torch.zeros(channels))
			self.register_buffer('exp_var',torch.ones(channels))
		
	def forward(self,x:torch.Tensor):
		#获取输入张量的形状，并确保通道数与 channels 参数相同
		x_shape=x.shape
		batch_size=x_shape[0]
		assert self.channels==x.shape[1]
		x=x.view(batch_size,self.channels,-1)
		if self.training or not self.track_running_stats:
			#计算输入张量 x 在第 0 和第 2 个维度上的均值
			mean=x.mean(dim=[0,2])
			mean_x2=(x**2).mean(dim=[0,2])
			var=mean_x2-mean**2
			#在训练过程中更新指数移动平均的均值和方差
			if self.training and self.track_running_stats:
				self.exp_mean=(1-self.momentum)*self.exp_mean+self.momentum*mean
				self.exp_var=(1-self.momentum)*self.exp_var+self.momentum*var
		else:
			mean=self.exp_mean
			var=self.exp_var
		#归一化
		x_norm=(x-mean.view(1,-1,1))/torch.sqrt(var+self.eps).view(1,-1,1)
		#拉伸和偏移
		if self.affine:
			x_norm=self.scale.view(1,-1,1)*x_norm+self.shift.view(1,-1,1)
		return x_norm.view(x_shape)

在批标准化中，拉伸（scale）和偏移（shift）是可学习的参数，它们的作用如下：

• 拉伸（Scale）：拉伸参数通常用一个可学习的缩放因子来表示，它乘以标准化后的特征值，用于调整特征的尺度。通过拉伸参数，批标准化能够在保持数据的零均值和单位方差的同时，对每个特征进行适当的缩放，从而提高网络的表达能力。如果数据在某个特征上有较大的方差，拉伸参数可以将其缩小；反之，如果数据在某个特征上有较小的方差，拉伸参数可以将其放大。
• 偏移（Shift）：偏移参数通常用一个可学习的偏移量来表示，它加到标准化后的特征值上，用于调整特征的平移。通过偏移参数，批标准化能够在保持数据的零均值和单位方差的同时，对每个特征进行适当的平移，从而使得网络能够学习到适合当前任务的特征表示。偏移参数允许网络自由地学习每个特征的偏移量，从而提高模型的灵活性。

综合来说，拉伸和偏移参数允许网络自适应地调整标准化后的特征，使得网络能够更好地适应不同的数据分布和任务需求，从而提高模型的性能和泛化能力。

参考：https://nn.labml.ai/normalization/batch_norm/index.html