pytorch bn层使用

在 PyTorch 中，Batch Normalization（BN）层可以通过 `torch.nn.BatchNorm` 模块来实现。BN层可以用于加速深度神经网络的收敛，并减少过拟合的风险。

下面是一个简单的例子，展示如何在 PyTorch 中使用 BN 层：

import torch
import torch.nn as nn

input_shape = (32, 3, 224, 224)

# 创建一个4D输入张量
x = torch.randn(input_shape)

# 创建一个包含Batch Normalization层的卷积神经网络
model = nn.Sequential(
    nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
    nn.BatchNorm2d(64),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
    nn.BatchNorm2d(128),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
    nn.BatchNorm2d(256),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Flatten(),
    nn.Linear(256 * 28 * 28, 1024),
    nn.BatchNorm1d(1024),
    nn.ReLU(),
    nn.Linear(1024, 10),
    nn.Softmax(dim=1)
)

# 打印模型结构
print(model)

# 前向计算
y = model(x)

在这个例子中，我们创建了一个包含 BN 层的简单的卷积神经网络。我们使用 `nn.BatchNorm2d` 和 `nn.BatchNorm1d` 分别添加了2D和1D的BN层。在前向计算中，我们将输入张量 `x` 传递给模型，得到输出张量 `y`。