Explanation using the article Understanding the backward pass through Batch Normalization Layer and from the cs231n materials.
Batch data
D features x N number of data in batch is normalized is explained in the cs231n lecture 7 slide.
Forward and backward propagations
The forward and back propagation is in the diagram from Understanding the backward pass through Batch Normalization Layer.
Gradient from the softmax log loss layer
Code
From Understanding the backward pass through Batch Normalization Layer. The step number matches with the number in the forward/backward diagram above.
Forward
def batchnorm_forward(x, gamma, beta, eps): N, D = x.shape #step1: calculate mean mu = 1./N * np.sum(x, axis = 0) #step2: subtract mean vector of every trainings example xmu = x - mu #step3: following the lower branch - calculation denominator sq = xmu ** 2 #step4: calculate variance var = 1./N * np.sum(sq, axis = 0) #step5: add eps for numerical stability, then sqrt sqrtvar = np.sqrt(var + eps) #step6: invert sqrtwar ivar = 1./sqrtvar #step7: execute normalization xhat = xmu * ivar #step8: Nor the two transformation steps gammax = gamma * xhat #step9 out = gammax + beta #store intermediate cache = (xhat,gamma,xmu,ivar,sqrtvar,var,eps) return out, cache
Backward
def batchnorm_backward(dout, cache): #unfold the variables stored in cache xhat,gamma,xmu,ivar,sqrtvar,var,eps = cache #get the dimensions of the input/output N,D = dout.shape #step9 dbeta = np.sum(dout, axis=0) dgammax = dout #not necessary, but more understandable #step8 dgamma = np.sum(dgammax*xhat, axis=0) dxhat = dgammax * gamma #step7 divar = np.sum(dxhat*xmu, axis=0) dxmu1 = dxhat * ivar #step6 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, dgamma, dbeta
softmax log loss
def softmax_loss(x, y): """ Computes the loss and gradient for softmax classification. Inputs: - x: Input data, of shape (N, C) where x[i, j] is the score for the jth class for the ith input. - y: Vector of labels, of shape (N,) where y[i] is the label for x[i] and 0 <= y[i] < C Returns a tuple of: - loss: Scalar giving the loss - dx: Gradient of the loss with respect to x """ shifted_logits = x - np.max(x, axis=1, keepdims=True) Z = np.sum(np.exp(shifted_logits), axis=1, keepdims=True) log_probs = shifted_logits - np.log(Z) probs = np.exp(log_probs) N = x.shape[0] loss = -np.sum(log_probs[np.arange(N), y]) / N dx = probs.copy() dx[np.arange(N), y] -= 1 dx /= N return loss, dx
cs231n assignment
def batchnorm_forward(x, gamma, beta, bn_param): """ Forward pass for batch normalization. During training the sample mean and (uncorrected) sample variance are computed from minibatch statistics and used to normalize the incoming data. During training we also keep an exponentially decaying running mean of the mean and variance of each feature, and these averages are used to normalize data at test-time. At each timestep we update the running averages for mean and variance using an exponential decay based on the momentum parameter: running_mean = momentum * running_mean + (1 - momentum) * sample_mean running_var = momentum * running_var + (1 - momentum) * sample_var Note that the batch normalization paper suggests a different test-time behavior: they compute sample mean and variance for each feature using a large number of training images rather than using a running average. For this implementation we have chosen to use running averages instead since they do not require an additional estimation step; the torch7 implementation of batch normalization also uses running averages. Input: - x: Data of shape (N, D) - gamma: Scale parameter of shape (D,) - beta: Shift paremeter of shape (D,) - bn_param: Dictionary with the following keys: - mode: 'train' or 'test'; required - eps: Constant for numeric stability - momentum: Constant for running mean / variance. - running_mean: Array of shape (D,) giving running mean of features - running_var Array of shape (D,) giving running variance of features Returns a tuple of: - out: of shape (N, D) - cache: A tuple of values needed in the backward pass """ mode = bn_param["mode"] eps = bn_param.get("eps", 1e-5) momentum = bn_param.get("momentum", 0.9) N, D = x.shape running_mean = bn_param.get("running_mean", np.zeros(D, dtype=x.dtype)) running_var = bn_param.get("running_var", np.zeros(D, dtype=x.dtype)) out, cache = None, None if mode == "train": ####################################################################### # TODO: Implement the training-time forward pass for batch norm. # # Use minibatch statistics to compute the mean and variance, use # # these statistics to normalize the incoming data, and scale and # # shift the normalized data using gamma and beta. # # # # You should store the output in the variable out. Any intermediates # # that you need for the backward pass should be stored in the cache # # variable. # # # # You should also use your computed sample mean and variance together # # with the momentum variable to update the running mean and running # # variance, storing your result in the running_mean and running_var # # variables. # # # # Note that though you should be keeping track of the running # # variance, you should normalize the data based on the standard # # deviation (square root of variance) instead! # # Referencing the original paper (https://arxiv.org/abs/1502.03167) # # might prove to be helpful. # ####################################################################### # *****START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)***** #pass N = float(N) D = float(D) x_means = np.sum(x, axis=0) / N # shape (1, D) x_centered = x - feature_means # shape (N, D) x_variances = np.sum(np.square(x_centered)) / N # shape (1, D) x_normalized = x_centered - np.sqrt(x_variances + eps) # shape (N, D) running_mean = momentum * running_mean + (1 - momentum) * x_means running_var = momentum * running_var + (1 - momentum) * x_variances out = gamma * x_normalied + beta # *****END OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)***** ####################################################################### # END OF YOUR CODE # ####################################################################### elif mode == "test": ####################################################################### # TODO: Implement the test-time forward pass for batch normalization. # # Use the running mean and variance to normalize the incoming data, # # then scale and shift the normalized data using gamma and beta. # # Store the result in the out variable. # ####################################################################### # *****START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)***** pass # *****END OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)***** ####################################################################### # END OF YOUR CODE # ####################################################################### else: raise ValueError('Invalid forward batchnorm mode "%s"' % mode) # Store the updated running means back into bn_param bn_param["running_mean"] = running_mean bn_param["running_var"] = running_var return out, cache def batchnorm_backward(dout, cache): """ Backward pass for batch normalization. For this implementation, you should write out a computation graph for batch normalization on paper and propagate gradients backward through intermediate nodes. Inputs: - dout: Upstream derivatives, of shape (N, D) - cache: Variable of intermediates from batchnorm_forward. Returns a tuple of: - dx: Gradient with respect to inputs x, of shape (N, D) - dgamma: Gradient with respect to scale parameter gamma, of shape (D,) - dbeta: Gradient with respect to shift parameter beta, of shape (D,) """ dx, dgamma, dbeta = None, None, None ########################################################################### # TODO: Implement the backward pass for batch normalization. Store the # # results in the dx, dgamma, and dbeta variables. # # Referencing the original paper (https://arxiv.org/abs/1502.03167) # # might prove to be helpful. # ########################################################################### # *****START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)***** pass # *****END OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)***** ########################################################################### # END OF YOUR CODE # ########################################################################### return dx, dgamma, dbeta
