Set track_running_stats=False during training,but running_mean still update

I want to fix the running_mean and running_var in BN during some training iteration.
So I set track_running_stats=False during training,but why running_mean still update.

bn = nn.BatchNorm2d(3)
for _ in range(10):
    x = torch.randn(10, 3, 24, 24)
    out = bn(x)

print(bn.running_mean)
> tensor([ 0.0031, -0.0034, -0.0066])
print(bn.running_var)
> tensor([1.0027, 0.9977, 1.0081])

print(bn)
> BatchNorm2d(3, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
bn.track_running_stats=False
print(bn)
> BatchNorm2d(3, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)

for _ in range(10):
    x = torch.randn(10, 3, 24, 24)
    out = bn(x)
print(bn.running_mean)
> tensor([ 0.0014, -0.0021, -0.0048])
print(bn.running_var)
> tensor([1.0036, 0.9980, 0.9983])

track_running_stats will be used to initialize the running stats in the __init__ method as seen here.
If you don’t want to update these stats, you should call bn.eval() and reset it afterwards via bn.train() to train them further.

Can you give an example code to illustrate the sequence of the process that you mentioned?

I am having the problem of using Batch normalization. The performance of my model drops significantly when running at inference mode (net.eval()).

I am looking forward to how to properly train and run the inference my model with Batch normalization with track_running_stats=False. Could you give the guidance.

If you set track_running_stats=False, the batchnorm layer will always use the batch statistics to normalize the activations, so there is no need to call train() or eval() on this layer anymore.

1 Like

I see. Thank you very much. :slight_smile: