Training with BatchNorm in pytorch
I'm wondering if I need to do anything special when training with BatchNorm in pytorch. From my understanding the gamma and beta parameters are updated with gradients as would normally be done by an optimizer. However, the mean and variance of the batches are updated slowly using momentum.


So do we need to specify to the optimizer when the mean and variance parameters are updated, or does pytorch automatically take care of this?
Is there a way to access the mean and variance of the BN layer so that I can make sure it was changing while I trained the model.


If needed here is my model and training procedure:

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)


Training:

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()


Batchnorm layers behave differently depending on if the model is in train or eval mode.
When net is in train mode (i.e. after calling net.train()) the batch norm layers contained in net will use batch statistics along with gamma and beta parameters to scale and translate each mini-batch. The running mean and variance will also be adjusted while in train mode. These updates to running mean and variance occur during the forward pass (when net(inputs) is called). The gamma and beta parameters are like any other pytorch parameter and are updated only once optimizer.step() is called.
When net is in eval mode (net.eval()) batch norm uses the historical running mean and running variance computed during training to scale and translate samples.
You can check the batch norm layers running mean and variance by displaying the layers running_mean and running_var members to ensure batch norm is updating them as expected. The learnable gamma and beta parameters can be accessed by displaying the weight and bias members of a batch norm layer respectively.
Edit
Below is a simple demonstration code showing that running_mean is updated during forward. Observe that it is not updated by the optimizer.
&gt;&gt;&gt; import torch
&gt;&gt;&gt; import torch.nn as nn
&gt;&gt;&gt; layer = nn.BatchNorm1d(5)
&gt;&gt;&gt; layer.train()
&gt;&gt;&gt; layer.running_mean
tensor([0., 0., 0., 0., 0.])
&gt;&gt;&gt; result = layer(torch.randn(5,5))
&gt;&gt;&gt; layer.running_mean
tensor([ 0.0271,  0.0152, -0.0403, -0.0703, -0.0056])

