title: Batch normalization in tensorflow: variables and performance
body: <p>I would like to add conditional operations on the variables of a batch normalization layer. Specifically, train in float, then quantize in a fine-tuning secondary training phase. For this, I want to add a tf.cond operation on the variables (scale, shift and exp moving averages of mean and var).</p>
  <p>I replaced the <code>tf.layers.batch_normalization</code>  with a batchnorm layer I wrote (see below).</p>
  <p>This function works perfectly (i.e. I get the same metrics with both functions), and I can add whatever pipeline to the variables (before the batchnorm operation). <strong>The problem is that the performance (runtime) dropped dramatically</strong> (i.e. there's a x2 factor by simply replacing the layers.batchnorm with my own function, as written below).</p>
  <pre><code>def batchnorm(self, x, name, epsilon=0.001, decay=0.99):
      epsilon = tf.to_float(epsilon)
      decay = tf.to_float(decay)
      with tf.variable_scope(name):
          shape = x.get_shape().as_list()
          channels_num = shape[3]
          # scale factor
          gamma = tf.get_variable(&quot;gamma&quot;, shape=[channels_num], initializer=tf.constant_initializer(1.0), trainable=True)
          # shift value
          beta = tf.get_variable(&quot;beta&quot;, shape=[channels_num], initializer=tf.constant_initializer(0.0), trainable=True)
          moving_mean = tf.get_variable(&quot;moving_mean&quot;, channels_num, initializer=tf.constant_initializer(0.0), trainable=False)
          moving_var = tf.get_variable(&quot;moving_var&quot;, channels_num, initializer=tf.constant_initializer(1.0), trainable=False)
          batch_mean, batch_var = tf.nn.moments(x, axes=[0, 1, 2]) # per channel
  
          update_mean = moving_mean.assign((decay * moving_mean) + ((1. - decay) * batch_mean))
          update_var = moving_var.assign((decay * moving_var) + ((1. - decay) * batch_var))
  
          tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, update_mean)
          tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, update_var)
  
          bn_mean = tf.cond(self.is_training, lambda: tf.identity(batch_mean), lambda: tf.identity(moving_mean))
          bn_var = tf.cond(self.is_training, lambda: tf.identity(batch_var), lambda: tf.identity(moving_var))
  
          with tf.variable_scope(name + &quot;_batchnorm_op&quot;):
              inv = tf.math.rsqrt(bn_var + epsilon)
              inv *= gamma
              output = ((x*inv) - (bn_mean*inv)) + beta
  
      return output
  </code></pre>
  <p>I would appreciate help in any of the following questions:</p>
  <ul>
  <li>Any ideas on how to improve the performance (reduce runtime) of my solution?</li>
  <li>Is it possible to add my own operators to the variables pipeline of layers.batchnorm before the batchnorm operation?</li>
  <li>Any other solution to the same problem?</li>
  </ul>