CNN中的批量归一化：原理剖析与新手实践指南

1次阅读

没有评论

共计 4090 个字符，预计需要花费 11 分钟才能阅读完成。

在训练深度神经网络时，尤其是 CNN，我们经常会遇到两个棘手的问题：

内部协变量偏移（Internal Covariate Shift）：随着网络层数的加深，每一层的输入分布会不断变化，导致后续层需要不断适应新的数据分布，这会显著减慢训练速度。
梯度消失 / 爆炸：在深层网络中，梯度在反向传播时可能会变得非常小或非常大，使得网络难以训练。

批量归一化（Batch Normalization，简称 BN）就是为了解决这些问题而提出的。它通过对每一层的输入进行归一化，使得网络各层的输入分布保持稳定，从而加速训练并提高模型性能。

BN 的核心思想很简单：对每个 mini-batch 的数据进行归一化，使其均值为 0，方差为 1。具体来说，对于输入特征 x，BN 的计算过程如下：

计算 mini-batch 的均值：
$$\mu_B = \frac{1}{m} \sum_{i=1}^m x_i$$
计算 mini-batch 的方差：
$$\sigma_B^2 = \frac{1}{m} \sum_{i=1}^m (x_i – \mu_B)^2$$
归一化：
$$\hat{x}_i = \frac{x_i – \mu_B}{\sqrt{\sigma_B^2 + \epsilon}}$$
缩放和平移（引入可学习参数 γ 和 β）：
$$y_i = \gamma \hat{x}_i + \beta$$

其中，ε 是一个很小的常数（如 1e-5），用于数值稳定性。γ 和 β 是可学习的参数，允许网络在需要时恢复原始数据分布。

在 PyTorch 中，我们可以使用 nn.BatchNorm2d 来实现 BN。例如：

import torch.nn as nn

# 定义一个带有 BN 的 CNN 层
conv_layer = nn.Sequential(nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, padding=1),
    nn.BatchNorm2d(64),
    nn.ReLU(),
    nn.MaxPool2d(2)
)

在 TensorFlow 中，对应的实现是 BatchNormalization 层：

from tensorflow.keras.layers import BatchNormalization, Conv2D, ReLU, MaxPool2D

# 定义一个带有 BN 的 CNN 层
model = tf.keras.Sequential([Conv2D(64, 3, padding='same'),
    BatchNormalization(),
    ReLU(),
    MaxPool2D()])

两者的 API 非常相似，但 PyTorch 的 BatchNorm2d 需要明确指定特征数（即通道数），而 TensorFlow 的 BatchNormalization 不需要。

下面我们用一个完整的 CNN 示例来展示 BN 的效果。我们将对比有无 BN 时的训练曲线。

import torch
import torchvision
import torchvision.transforms as transforms

# 数据预处理
transform = transforms.Compose([transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# 加载 CIFAR-10 数据集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True)

# 定义有 BN 和没有 BN 的两个模型
class CNNWithBN(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Conv2d(3, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(64*8*8, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Linear(256, 10)
        )

    def forward(self, x):
        return self.net(x)

class CNNWithoutBN(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Conv2d(3, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(64*8*8, 256),
            nn.ReLU(),
            nn.Linear(256, 10)
        )

    def forward(self, x):
        return self.net(x)

import torch.optim as optim

# 初始化模型
model_with_bn = CNNWithBN().cuda()
model_without_bn = CNNWithoutBN().cuda()

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer_bn = optim.Adam(model_with_bn.parameters(), lr=0.001)
optimizer_no_bn = optim.Adam(model_without_bn.parameters(), lr=0.001)

# 训练函数
def train(model, optimizer, trainloader, epochs=10):
    model.train()
    losses = []
    for epoch in range(epochs):
        running_loss = 0.0
        for i, (inputs, labels) in enumerate(trainloader):
            inputs, labels = inputs.cuda(), labels.cuda()

            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            if i % 100 == 99:
                print(f'Epoch {epoch+1}, Batch {i+1}, Loss: {running_loss/100:.3f}')
                losses.append(running_loss/100)
                running_loss = 0.0
    return losses

# 训练两个模型
losses_with_bn = train(model_with_bn, optimizer_bn, trainloader)
losses_without_bn = train(model_without_bn, optimizer_no_bn, trainloader)

图：有 BN 和无 BN 时的训练损失对比