Как тензорный поток вычисляет градиенты через слой пакетной нормализации?

Question

Мне нужно повторить вычисления, которые тензорный поток делает при распространении обратно через слой нормализации партии. Я использую эту реализацию batchnorm в качестве шаблона:

class batch_norm:
def __init__(self,inputs,size,is_training,sess,bn_param=None):

    self.sess = sess        
    self.scale = tf.Variable(tf.random_uniform([size],1.,1.))
    self.beta = tf.Variable(tf.random_uniform([size],0.,0.))
    self.pop_mean = tf.Variable(tf.random_uniform([size],0.,0.)) #used when in testing mode
    self.pop_var = tf.Variable(tf.random_uniform([size],1.,1.))  #used when in testing mode     
    self.batch_mean, self.batch_var = tf.nn.moments(inputs,[0])        
    self.train_mean = tf.assign(self.pop_mean,self.pop_mean * decay + self.batch_mean * (1 - decay))  
    self.train_var = tf.assign(self.pop_var,self.pop_var * decay + self.batch_var * (1 - decay))

    def training(): 
        return (inputs - self.batch_mean) / tf.sqrt(self.batch_var + 0.0000001) * self.scale + self.beta

    def testing(): 
        return (inputs - self.pop_mean) / tf.sqrt(self.pop_var + 0.0000001) * self.scale + self.beta

Мне нужна наивная реализация вычислений, которые тензорный поток делает при распространении обратно через него. Я взял реализацию batchnorm, описанную в этой статье, и провел простой эксперимент со слоем из одного нейрона, передавая наивной и tf-версии одинаковые значения, и результирующие градиенты сильно отличаются. Это наивная реализация обратного прохода:

 def batchnorm_backward_training(self, dout, x, gamma):

    #get the dimensions of the input/output
    N,D = dout.shape

    mu = 1./N * np.sum(x, axis = 0)  #calculate the mean

    xmu = x - mu  #substract the mean from the input

    sq = xmu ** 2

    var = 1./N * np.sum(sq, axis = 0)  #calculate the variance

    eps = 0.0000001

    sqrtvar = np.sqrt(var + eps)

    ivar = 1./sqrtvar

    xhat = xmu * ivar  #normalized input


    dbeta = np.sum(dout, axis=0)  #gradients for beta
    dgammax = dout   #gradients for gamma

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

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

    dsqrtvar = -1. /(sqrtvar**2) * divar

    #step5
    dvar = 0.5 * 1. /np.sqrt(var+eps) * dsqrtvar

    #step4
    dsq = 1. /N * np.ones((N,D)) * dvar

    #step3
    dxmu2 = 2 * xmu * dsq

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

    #step1
    dx2 = 1. /N * np.ones((N,D)) * dmu

    #step0
    dx = dx1 + dx2

    return dx  #final gradient

Градиенты в эксперименте варьируются как во время обучения, так и во время тестирования (разница между тестированием и обучением для наивной реализации заключается в том, что дисперсия и среднее значение устанавливаются равными 1 и 0 во время тестирования). Вот ссылка на полный код эксперимента. Из того, что я мог найти в Интернете, все другие реализации обратного распространения batchnorm используют некоторую версию кода, представленную в статье.

Как мне нужно изменить его, чтобы получить правильный результат?