title: What are ways I can debug this keras layer?
body: <p>I am new to  Keras and am attempting to implement the decorrelated batch Norm paper (<a href="https://arxiv.org/abs/1804.08450" rel="nofollow noreferrer">https://arxiv.org/abs/1804.08450</a>) in Keras as a learning experience. The layer is very similar to standard batch norm with a few additional components.</p>
  
  <p>Instead of centering the input data to each layer and normalizing by the variance, we now center the data and apply a whitening transform which is computed by doing an eigenvalue decomposition on the covariance matrix. </p>
  
  <p>The entire procedure is clearly laid out in the paper (Algorithm 1, page 5) and composed of only like 5 equations, whose implementation I marked in the code below. I successfully re-implemented standard batch norm layer, but am getting NaN loss and low accuracy when I incorporate the whitening procedure. </p>
  
  <p>I am wondering if there any advice I should follow to debug this code. I am not sure if I made a dimensionality mistake or incorrectly implemented the equations, but any help would be appreciated.</p>
  
  <p>Here is the code if you are interested (Edited to include Daniel Möller's corrections). The input to the layer is a tensor of dimension (batch_size height width channels).</p>
  
  <pre><code>input_shape = K.int_shape(inputs) # (batch_size height width channels) 
  # unroll all dimensions except feature maps dim (c X hwb)
  pool_shape = (-1, input_shape[-1]) 
  x = K.reshape(x,pool_shape)
  x = K.permute_dimensions(x, (1,0)) #if you do want to invert the dimensions
  
  mean = K.mean(x,1,keepdims=True)     
  
  # standard batch norm
  #stddev = K.std(x,1,keepdims=True) + self.epsilon
  #normed = (x - mean) / stddev
  #normed = K.reshape(normed,((-1,)+ input_shape[1:]))
  
  # center inputs
  centered_inputs = x - mean 
  
  #vvvvvERROR SOMEWHERE IN HEREvvvvv#
  # compute covariance matrix for reshaped inputs xxt
  covar = K.batch_dot(K.expand_dims(x, axis=-1), K.expand_dims(x, axis=-1),axes=(2,2))
  # fuzz covariance matrix to prevent singularity
  covar = covar + self.epsilon 
  
  # execute eigenvalue decomposition
  #Lambda, D,_ = tf.svd(covar,compute_uv=True)
  Lambda, D = tf.self_adjoint_eig(covar)
  Lambda = tf.linalg.diag(Lambda)
  
  # calculate PCA-whitening matrix 1/sqrt(L) * D^T
  U = K.batch_dot(1. / K.sqrt(Lambda), D, axes=(2,2))
  # calculate PCA-whitened activation x_a = U(x - \mu)
  x_a = K.batch_dot(U, centered_inputs,axes=(2,1))
  # calculate ZCA-whitened output Dx_a
  x_whitened = K.batch_dot(D, x_a)
  #^^^^^ERROR SOMEWHERE IN HERE^^^^^# 
  
  # reshape whitened activations back to input dimension
  x_normed = K.permute_dimensions(x_whitened,(1,0)) # permute back to (bhw X c)
  x_normed = K.reshape(x_normed,((-1,), input_shape[1:])) # reroll dimensions
  </code></pre>