PyTorch中怎么实现批量归一化

import torch
import torch.nn as nn

# 对输入数据进行批量归一化
input_data = torch.randn(20, 16, 50, 50)  # 输入数据的shape为(batch_size, channels, height, width)

# 对2D数据进行批量归一化
batchnorm = nn.BatchNorm2d(16)  # 对通道维度进行批量归一化
output_data = batchnorm(input_data)

# 对1D数据进行批量归一化
input_data = torch.randn(20, 16, 100)  # 输入数据的shape为(batch_size, channels, length)
batchnorm = nn.BatchNorm1d(16)  # 对特征维度进行批量归一化
output_data = batchnorm(input_data)