在训练神经网络时,为了避免梯度爆炸或梯度消失的问题,通常会对网络进行归一化操作。Batch Normalization(批量归一化)是一种有效的归一化操作方法,它可以在数据被送入激活函数之前,将数据进行归一化,使得数据的分布更加稳定,有利于网络的训练。
在PyTorch中,使用Batch Normalization层需要使用model.train()来表示进入训练模式,使用model.eval()来表示进入评估模式。在评估模式下,Batch Normalization层将使用前面训练得到的均值和方差来对数据进行归一化,而不是对当前的mini-batch进行均值和方差的计算,以保证测试时得到的结果的稳定性。
代码示例:
import torch import torch.nn as nn
input = torch.randn(32, 3, 64, 64) model = nn.Sequential( nn.Conv2d(3, 16, 3), nn.BatchNorm2d(16), # Batch Normalization层 nn.ReLU(), nn.Conv2d(16, 32, 3), nn.BatchNorm2d(32), # Batch Normalization层 nn.ReLU(), nn.MaxPool2d(2), nn.Flatten(), nn.Linear(321515, 10), )
model.train() output_train = model(input)
model.eval() output_eval = model(input)