The differences of BatchNorm layer backpropagation at mode of model.train() and model.eval() in Pytorch?

Question

I test the gradient of BatchNorm layer for two mode: model.train() and model.eval(). I bulid a simple CNN network NetWork and input the same input X to the network at model.train() mode and model.eval() mode. I know the differences of model.train() and model.eval() of BatchNorm layer. I have replaced the mean and the var of Batchnorm layer in mode of model.eval() as the values in mode of model.train(). Therefore both outputs and parameters of two mode are the same. output of two mode parameters of two mode However, when I calculate the gradients of each parameters, I found the gradients of the layer before BatchNorm layer is different, although their parameters and the loss is same. different gradients of the layer before BatchNorm I think it's because the difference of BatchNorm layer backpropagation at model.train() and model.eval(), but I don't understrand the detail of it. Does anyone know? Thank you so much.

Answer 1

When the mode is.train(), the batchnorm layer calculate the batchwise mean and variance of the input and uses it to normalize the inputs. This mean and variance is also used to update the Moving Average Mean and Variance.

When the mode is.eval(), the batchnorm layer doesn't calculate the mean and variance of the input, but uses the pre-computed moving average mean and variance during training stage.

This way, your predictions won't change on a single image during testing, when other samples in the batch changes.

https://pytorch.org/docs/stable/_modules/torch/nn/modules/batchnorm.html#BatchNorm2d

In the above code, running mean and running variance is the moving average mean and variance of the input featuremap of batch norm layer that they were calculated during training stage.