batchNormlize torch 使用

BatchNormlize 是一种常用的正则化方式，在 PyTorch 中可以通过 torch.nn.BatchNorm1d 或 torch.nn.BatchNorm2d 模块来实现。这两个模块分别用于对一维和二维的输入数据进行批量归一化。以下是一个简单的示例代码：

import torch
import torch.nn as nn

# 定义一个简单的神经网络
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
        self.bn1 = nn.BatchNorm2d(6)   # BatchNorm2d 模块

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)   # 批量归一化
        x = torch.relu(x)
        x = self.pool(x)
        x = self.conv2(x)
        x = torch.relu(x)
        x = self.pool(x)
        x = x.view(-1, 16 * 5 * 5)
        x = self.fc1(x)
        x = torch.relu(x)
        x = self.fc2(x)
        x = torch.relu(x)
        x = self.fc3(x)
        return x

# 使用批量归一化的神经网络进行训练
net = Net()
inputs = torch.randn(1, 3, 32, 32)
outputs = net(inputs)

在上面的代码中，我们定义了一个简单的神经网络，并在第一个卷积层后添加了 BatchNorm2d 模块，用于对输入数据进行批量归一化。在 forward 函数中，我们调用了 BatchNorm2d 模块，将卷积层的输出进行归一化处理。这样可以加速模型的训练，并提高模型的准确率。