PyTorch Batch Normalization Guide

In PyTorch, batch normalization layers can be implemented using the BatchNorm1d, BatchNorm2d, or BatchNorm3d classes within the torch.nn module. These classes are used to apply batch normalization on 1D, 2D, or 3D data respectively.

Here is a simple example demonstrating how to use batch normalization in PyTorch.

import torch
import torch.nn as nn

# 创建一个简单的神经网络模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init()
        self.fc1 = nn.Linear(10, 20)
        self.bn1 = nn.BatchNorm1d(20)
        self.fc2 = nn.Linear(20, 10)
        self.bn2 = nn.BatchNorm1d(10)
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.bn1(x)
        x = nn.ReLU(x)
        x = self.fc2(x)
        x = self.bn2(x)
        x = nn.ReLU(x)
        return x

# 初始化模型
model = Net()

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

# 训练模型
for epoch in range(10):
    for data, target in train_loader:
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

In the code above, we have constructed a basic neural network model with batch normalization layers. We have then specified the loss function and optimizer, and trained the model using the data in the train_loader.

Note that we have applied batch normalization layers in the forward() method of the model. This ensures that during the training process, each batch of input data is normalized, speeding up training and improving model performance.