在 pytorch 中使用 BatchNorm 进行训练

Question

我想知道在 pytorch 中使用 BatchNorm 进行训练时是否需要做一些特别的事情。根据我的理解，gamma和beta参数使用梯度更新，这通常由优化器完成。然而，批次的均值和方差使用动量缓慢更新。

1. 那么我们需要在均值和方差参数更新时指定给优化器，还是 pytorch 会自动处理这个问题？
2. 有没有办法访问 BN 层的均值和方差，以便我可以确保在训练模型时它正在发生变化。

如果需要，这里是我的模型和训练程序：

def bn_drop_lin(n_in:int, n_out:int, bn:bool=True, p:float=0.):
    "Sequence of batchnorm (if `bn`), dropout (with `p`) and linear (`n_in`,`n_out`) layers followed by `actn`."
    layers = [nn.BatchNorm1d(n_in)] if bn else []
    if p != 0: layers.append(nn.Dropout(p))
    layers.append(nn.Linear(n_in, n_out))

    return nn.Sequential(*layers)

class Model(nn.Module):
    def __init__(self, i, o, h=()):
        super().__init__()

        nodes = (i,) + h + (o,)
        self.layers = nn.ModuleList([bn_drop_lin(i,o, p=0.5) 
                                     for i, o in zip(nodes[:-1], nodes[1:])])

    def forward(self, x):
        x = x.view(x.shape[0], -1)
        for layer in self.layers[:-1]:
            x = F.relu(layer(x))

        return self.layers[-1](x)

训练：

for i, data in enumerate(trainloader):
    # get the inputs; data is a list of [inputs, labels]
    inputs, labels = data

    # zero the parameter gradients
    optimizer.zero_grad()

    # forward + backward + optimize
    outputs = net(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()

Answer 1

Batchnorm 层的行为取决于模型是处于训练模式还是评估模式。

当net处于训练模式时（即调用之后net.train()），其中包含的批规范层net将使用批统计数据以及 gamma 和 beta 参数来缩放和转换每个小批量。在训练模式下，运行均值和方差也将被调整。这些对运行均值和方差的更新发生在前向传递期间（net(inputs)调用时）。gamma 和 beta 参数就像任何其他 pytorch 参数一样，并且只更新一次optimizer.step()。

When netis in eval mode ( net.eval()) batch norm 使用训练期间计算的历史运行均值和运行方差来缩放和转换样本。

您可以通过显示层running_mean和running_var成员来检查运行均值和方差的批处理规范层，以确保批处理规范按预期更新它们。可以通过分别显示批处理规范层的weight和bias成员来访问可学习的 gamma 和 beta 参数。

编辑

下面是一个简单的演示代码，显示running_mean在转发期间更新。观察它没有被优化器更新。

>>> import torch
>>> import torch.nn as nn
>>> layer = nn.BatchNorm1d(5)
>>> layer.train()
>>> layer.running_mean
tensor([0., 0., 0., 0., 0.])
>>> result = layer(torch.randn(5,5))
>>> layer.running_mean
tensor([ 0.0271,  0.0152, -0.0403, -0.0703, -0.0056])