visualizing batch_norm parameters in tensorboard

Question

My current NN model is giving some anomalous results when I change batch norm specific hyper parameters. I'd like to see the distribution of the batch norm parameters beta and gamma over time to make sure that batch norm isn't doing something weird.

Visualizing learned weights or biases is easiest to do with tensorboard, but I'm not sure how to do that with beta and gamma since they're defined and managed within tf.layers.batch_normalization or tf.contrib.layers.batch_norm.

Is there a simple way to reference beta and gamma and put them in a histogram summary without having to write my own version of batch norm?

Answer 1

building a summary for them is still a chore, but this is what I've come up with for accessing gamma and beta:

def batch_norm(self, x_in):
    with tf.variable_scope('batch_norm'):
        x = tf.layers.batch_normalization(  x_in,
                                            momentum = self.bn_decay,
                                            epsilon = self.bn_epsilon,
                                            training = self.is_training)

        gamma =  tf.trainable_variables(tf.get_variable_scope().name)[0]
        beta = tf.trainable_variables(tf.get_variable_scope().name)[1]

        return x

what tf.trainable_variables(tf.get_variable_scope().name) does is return all variables within the current scope in the form of a list. In this case there are two variables, the 0th is gamma and the 1st is beta but that may change with a different implementation.

if you need the specific names use:

for var in tf.trainable_variables(tf.get_variable_scope().name):
    print(var.name)