Where is the source code for Batchnorm2d?

autograd
1

Hello,

recently, I tried to implement the batchnorm 2d layer with the scale layer.
the scale layer is doing like below
y = (scale * x + shift)^power

The batchnorm2d layer consists of two consecutive scale layers as I know.

  1. nomalization(zero mean and unit variance)
    hat(x) = (x-E[x]) / sqrt(Var[X] + eps)

  2. affine transformation.
    y = gamma * hat(x) + beta

I found the batchnorm2d has four parameters. running_mean, running_var, weight, bias.

So…I implemented like below

  1. the first layer do normalization
    scale = 1 / sqrt(running_var+eps)
    shift = running_mean * -1.0
    power = 1.0

  2. and second layer do the affine transformation
    scale = weights
    shift = bias
    power = 1.0

I think it looks well but I got the wrong values from my implementation.
So, I want to know how to operate the batchnorm2d in evaluation mode.
Where can I find?