如何在pytorch里lenet中添加BatchNormalization

要在LeNet中添加批量归一化（BatchNormalization），可以在每个卷积层或全连接层之后添加BatchNormalization层。这可以通过使用PyTorch中的nn.BatchNorm2d或nn.BatchNorm1d模块来实现。

以下是一个示例代码，展示如何在LeNet中添加批量归一化：

import torch.nn as nn

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.bn1 = nn.BatchNorm2d(6)  # BatchNormalization层
        self.pool1 = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.bn2 = nn.BatchNorm2d(16)  # BatchNormalization层
        self.pool2 = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.bn3 = nn.BatchNorm1d(120)  # BatchNormalization层
        self.fc2 = nn.Linear(120, 84)
        self.bn4 = nn.BatchNorm1d(84)  # BatchNormalization层
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)  # 添加BatchNormalization层
        x = nn.functional.relu(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)  # 添加BatchNormalization层
        x = nn.functional.relu(x)
        x = self.pool2(x)
        x = x.view(-1, 16 * 5 * 5)
        x = self.fc1(x)
        x = self.bn3(x)  # 添加BatchNormalization层
        x = nn.functional.relu(x)
        x = self.fc2(x)
        x = self.bn4(x)  # 添加BatchNormalization层
        x = nn.functional.relu(x)
        x = self.fc3(x)
        return x

在这个示例代码中，我们添加了4个BatchNormalization层，分别在conv1、conv2、fc1和fc2之后。注意，在fc1和fc2之后，我们使用了nn.BatchNorm1d模块，因为这些层是全连接层，而不是卷积层。

添加BatchNormalization层可以帮助提高模型的训练速度和准确性，因为它可以使得输入数据的分布更加稳定，从而加速训练过程