I would like to implement batch normalization as originally described in the paper (https://arxiv.org/pdf/1502.03167.pdf 17). After the network finishes training, they take one more pass through the dataset to estimate the population mean and variance (step 10 in algorithm 2). Pytorch instead uses the common trick of keeping a running average to estimate these statistics while training.
Is there an easy way to implement this? Basically I need to have access to the mini batch mean and variance computed within each batch normalization layer.
Is there code for this already?
Thanks!