coding utf8 references https kratzert github io20160212understandingthegradientflowthroughthebatchnormalizationlayer html step1 calculate mean step2 subtract mean vector of every trainings example step3 following the lower branch calculation denominator step4 calculate variance step5 add eps for numerical stability then sqrt step6 invert sqrtwar step7 execute normalization step8 nor the two transformation steps step9 store running averages of mean and variance during training for use during testing store intermediate input is a regular layer input is a convolution layer unfold the variables stored in cache get the dimensions of the inputoutput step9 step8 step7 step6 step5 step4 step3 step2 step1 step0 update gradient values input is a regular layer input is a convolution layer coding utf 8 references https kratzert github io 2016 02 12 understanding the gradient flow through the batch normalization layer html step1 calculate mean step2 subtract mean vector of every trainings example step3 following the lower branch calculation denominator step4 calculate variance step5 add eps for numerical stability then sqrt step6 invert sqrtwar step7 execute normalization step8 nor the two transformation steps step9 store running averages of mean and variance during training for use during testing store intermediate input is a regular layer input is a convolution layer unfold the variables stored in cache get the dimensions of the input output step9 not necessary but more understandable step8 step7 step6 step5 step4 step3 step2 step1 step0 update gradient values input is a regular layer input is a convolution layer
import numpy as np

from mla.neuralnet.layers import Layer, PhaseMixin, ParamMixin
from mla.neuralnet.parameters import Parameters




class BatchNormalization(Layer, ParamMixin, PhaseMixin):
    def __init__(self, momentum=0.9, eps=1e-5, parameters=None):
        super().__init__()
        self._params = parameters
        if self._params is None:
            self._params = Parameters()
        self.momentum = momentum
        self.eps = eps
        self.ema_mean = None
        self.ema_var = None

    def setup(self, x_shape):
        self._params.setup_weights((1, x_shape[1]))

    def _forward_pass(self, X):
        gamma = self._params["W"]
        beta = self._params["b"]

        if self.is_testing:
            mu = self.ema_mean
            xmu = X - mu
            var = self.ema_var
            sqrtvar = np.sqrt(var + self.eps)
            ivar = 1.0 / sqrtvar
            xhat = xmu * ivar
            gammax = gamma * xhat
            return gammax + beta

        N, D = X.shape

        
        mu = 1.0 / N * np.sum(X, axis=0)

        
        xmu = X - mu

        
        sq = xmu ** 2

        
        var = 1.0 / N * np.sum(sq, axis=0)

        
        sqrtvar = np.sqrt(var + self.eps)

        
        ivar = 1.0 / sqrtvar

        
        xhat = xmu * ivar

        
        gammax = gamma * xhat

        
        out = gammax + beta

        
        if self.ema_mean is None or self.ema_var is None:
            self.ema_mean = mu
            self.ema_var = var
        else:
            self.ema_mean = self.momentum * self.ema_mean + (1 - self.momentum) * mu
            self.ema_var = self.momentum * self.ema_var + (1 - self.momentum) * var
        
        self.cache = (xhat, gamma, xmu, ivar, sqrtvar, var)

        return out

    def forward_pass(self, X):
        if len(X.shape) == 2:
            
            return self._forward_pass(X)
        elif len(X.shape) == 4:
            
            N, C, H, W = X.shape
            x_flat = X.transpose(0, 2, 3, 1).reshape(-1, C)
            out_flat = self._forward_pass(x_flat)
            return out_flat.reshape(N, H, W, C).transpose(0, 3, 1, 2)
        else:
            raise NotImplementedError("Unknown model with dimensions = {}".format(len(X.shape)))

    def _backward_pass(self, delta):
        
        xhat, gamma, xmu, ivar, sqrtvar, var = self.cache

        
        N, D = delta.shape

        
        dbeta = np.sum(delta, axis=0)
        dgammax = delta  

        
        dgamma = np.sum(dgammax * xhat, axis=0)
        dxhat = dgammax * gamma

        
        divar = np.sum(dxhat * xmu, axis=0)
        dxmu1 = dxhat * ivar

        
        dsqrtvar = -1.0 / (sqrtvar ** 2) * divar

        
        dvar = 0.5 * 1.0 / np.sqrt(var + self.eps) * dsqrtvar

        
        dsq = 1.0 / N * np.ones((N, D)) * dvar

        
        dxmu2 = 2 * xmu * dsq

        
        dx1 = dxmu1 + dxmu2
        dmu = -1 * np.sum(dxmu1 + dxmu2, axis=0)

        
        dx2 = 1.0 / N * np.ones((N, D)) * dmu

        
        dx = dx1 + dx2

        
        self._params.update_grad("W", dgamma)
        self._params.update_grad("b", dbeta)

        return dx

    def backward_pass(self, X):
        if len(X.shape) == 2:
            
            return self._backward_pass(X)
        elif len(X.shape) == 4:
            
            N, C, H, W = X.shape
            x_flat = X.transpose(0, 2, 3, 1).reshape(-1, C)
            out_flat = self._backward_pass(x_flat)
            return out_flat.reshape(N, H, W, C).transpose(0, 3, 1, 2)
        else:
            raise NotImplementedError("Unknown model shape: {}".format(X.shape))

    def shape(self, x_shape):
        return x_shape