nn.BatchNorm2d 如何使用

nn.BatchNorm2d是PyTorch中的一个模块，用于对特征图进行批标准化，可以用于加速模型训练和提高模型精度。

使用步骤如下：

1. 导入模块：

import torch.nn as nn

2. 创建BatchNorm2d对象：

batch_norm = nn.BatchNorm2d(num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)

其中，num_features表示特征图的通道数，eps是一个小数，用于防止分母为0，momentum用于计算移动平均值的动量因子，affine表示是否要对BatchNorm的结果进行缩放和平移，track_running_stats表示是否要跟踪移动平均值和方差。

3. 将BatchNorm2d对象应用到网络中：

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu1 = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        
    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool(x)
        return x

在上面的代码中，我们将BatchNorm2d对象应用到了Conv2d之后，通过将它们串联起来，使得特征图在经过Conv2d之后直接进入BatchNorm2d，然后再进入ReLU激活函数。

4. 训练模型：

net = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
for epoch in range(10):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print('[%d] loss: %.3f' % (epoch + 1, running_loss / len(trainloader)))

在训练模型的过程中，我们可以使用标准的训练循环，通过反向传播和优化器来训练模型。

以上就是使用nn.BatchNorm2d的基本步骤。