Set track_running_stats=False during training，but running_mean still update

vision
#1

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])
#2

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.

#3

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.

#4

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
#5

I see. Thank you very much. :slight_smile: