Каковы способы отладки этого слоя keras?

Я новичок в Керасе и пытаюсь реализовать декоррелированную пакетную норму бумаги (https://arxiv.org/abs/1804.08450 ) в Керасе как опыт обучения. Слой очень похож на стандартную пакетную норму с несколькими дополнительными компонентами.

Вместо центрирования входных данных на каждом слое и нормализации по дисперсии, мы теперь центрируем данные и применяем преобразование отбеливания, которое вычисляется путем декомпозиции собственных значений на матрице ковариации.

Вся процедура четко изложена в статье (алгоритм 1, стр. 5) и состоит только из подобных 5 уравнений, реализация которых отмечена в приведенном ниже коде. Я успешно повторно внедрил стандартный пакетный слой norm, но получаю потерю наночастиц и низкую точность, когда я включаю процедуру отбеливания.

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

Вот код, если вам интересно (отредактировано, чтобы включить исправления Даниэля Меллера). Входной слой является тензорной размерности (batch_size высота ширина каналов).

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