How to change the mean and var value of the BN layer for a pretrained model?

I have a pretrained model, and want to apply it to a new dataset. From the AdaBN paper, I want to recaculate the the mean and var for each neuros and update the value for each BN layer. How could I reset the parameter?

You can call .reset_parameters() directly:

bn = nn.BatchNorm2d(3)
bn(torch.randn(10, 3, 24, 24))
print(bn.running_mean)
print(bn.running_var)
bn.reset_parameters()
print(bn.running_mean)
print(bn.running_var)
1 Like