(x,
       params=None,
       moments=None,
       backprop_through_moments=True,
       use_ema=False,
       is_training=True,
       ema_epsilon=.9)
Batch normalization.

  The usage should be as follows: If x is the support images, moments should be
  None so that they are computed from the support set examples. On the other
  hand, if x is the query images, the moments argument should be used in order
  to pass in the mean and var that were computed from the support set.

  Args:
    x: inputs.
    params: None or a dict containing the values of the offset and scale params.
    moments: None or a dict containing the values of the mean and var to use for
      batch normalization.
    backprop_through_moments: Whether to allow gradients to flow through the
      given support set moments. Only applies to non-transductive batch norm.
    use_ema: apply moving averages of batch norm statistics, or update them,
      depending on whether we are training or testing.  Note that passing
      moments will override this setting, and result in neither updating or
      using ema statistics.  This is important to make sure that episodic
      learners don't update ema statistics a second time when processing
      queries.
    is_training: if use_ema=True, this determines whether to apply the moving
      averages, or update them.
    ema_epsilon: if updating moving averages, use this value for the exponential
      moving averages.

  Returns:
    output: The result of applying batch normalization to the input.
    params: The updated params.
    moments: The updated moments.
Batch normalization.
Batch normalization .
def bn(x,
       params=None,
       moments=None,
       backprop_through_moments=True,
       use_ema=False,
       is_training=True,
       ema_epsilon=.9):
  """Batch normalization.

  The usage should be as follows: If x is the support images, moments should be
  None so that they are computed from the support set examples. On the other
  hand, if x is the query images, the moments argument should be used in order
  to pass in the mean and var that were computed from the support set.

  Args:
    x: inputs.
    params: None or a dict containing the values of the offset and scale params.
    moments: None or a dict containing the values of the mean and var to use for
      batch normalization.
    backprop_through_moments: Whether to allow gradients to flow through the
      given support set moments. Only applies to non-transductive batch norm.
    use_ema: apply moving averages of batch norm statistics, or update them,
      depending on whether we are training or testing.  Note that passing
      moments will override this setting, and result in neither updating or
      using ema statistics.  This is important to make sure that episodic
      learners don't update ema statistics a second time when processing
      queries.
    is_training: if use_ema=True, this determines whether to apply the moving
      averages, or update them.
    ema_epsilon: if updating moving averages, use this value for the exponential
      moving averages.

  Returns:
    output: The result of applying batch normalization to the input.
    params: The updated params.
    moments: The updated moments.
  """
  params_keys, params_vars, moments_keys, moments_vars = [], [], [], []

  with tf.variable_scope('batch_norm'):
    scope_name = tf.get_variable_scope().name

    if use_ema:
      ema_shape = [1, 1, 1, x.get_shape().as_list()[-1]]
      mean_ema = tf.get_variable(
          'mean_ema',
          shape=ema_shape,
          initializer=tf.initializers.zeros(),
          trainable=False)
      var_ema = tf.get_variable(
          'var_ema',
          shape=ema_shape,
          initializer=tf.initializers.ones(),
          trainable=False)

    if moments is not None:
      if backprop_through_moments:
        mean = moments[scope_name + '/mean']
        var = moments[scope_name + '/var']
      else:
        # This variant does not yield good resutls.
        mean = tf.stop_gradient(moments[scope_name + '/mean'])
        var = tf.stop_gradient(moments[scope_name + '/var'])
    elif use_ema and not is_training:
      mean = mean_ema
      var = var_ema
    else:
      # If not provided, compute the mean and var of the current batch.

      replica_ctx = tf.distribute.get_replica_context()
      if replica_ctx:
        # from tensorflow/python/keras/layers/normalization/batch_normalization.py  # pylint: disable=line-too-long
        axes = list(range(len(x.shape) - 1))
        local_sum = tf.reduce_sum(x, axis=axes, keepdims=True)
        local_squared_sum = tf.reduce_sum(
            tf.square(x), axis=axes, keepdims=True)
        batch_size = tf.cast(tf.shape(x)[0], tf.float32)
        x_sum, x_squared_sum, global_batch_size = (
            replica_ctx.all_reduce('sum',
                                   [local_sum, local_squared_sum, batch_size]))

        axes_vals = [(tf.shape(x))[i] for i in range(1, len(axes))]
        multiplier = tf.cast(tf.reduce_prod(axes_vals), tf.float32)
        multiplier = multiplier * global_batch_size

        mean = x_sum / multiplier
        x_squared_mean = x_squared_sum / multiplier
        # var = E(x^2) - E(x)^2
        var = x_squared_mean - tf.square(mean)
      else:
        mean, var = tf.nn.moments(
            x, axes=list(range(len(x.shape) - 1)), keep_dims=True)

    # Only update ema's if training and we computed the moments in the current
    # call.  Note: at test time for episodic learners, ema's may be passed
    # from the support set to the query set, even if it's not really needed.
    if use_ema and is_training and moments is None:
      replica_ctx = tf.distribute.get_replica_context()
      mean_upd = tf.assign(mean_ema,
                           mean_ema * ema_epsilon + mean * (1.0 - ema_epsilon))
      var_upd = tf.assign(var_ema,
                          var_ema * ema_epsilon + var * (1.0 - ema_epsilon))
      updates = tf.group([mean_upd, var_upd])
      if replica_ctx:
        tf.add_to_collection(
            tf.GraphKeys.UPDATE_OPS,
            tf.cond(
                tf.equal(replica_ctx.replica_id_in_sync_group, 0),
                lambda: updates, tf.no_op))
      else:
        tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, updates)

    moments_keys += [scope_name + '/mean']
    moments_vars += [mean]
    moments_keys += [scope_name + '/var']
    moments_vars += [var]

    if params is None:
      offset = tf.get_variable(
          'offset',
          shape=mean.get_shape().as_list(),
          initializer=tf.initializers.zeros())
      scale = tf.get_variable(
          'scale',
          shape=var.get_shape().as_list(),
          initializer=tf.initializers.ones())
    else:
      offset = params[scope_name + '/offset']
      scale = params[scope_name + '/scale']

    params_keys += [scope_name + '/offset']
    params_vars += [offset]
    params_keys += [scope_name + '/scale']
    params_vars += [scale]

    output = tf.nn.batch_normalization(x, mean, var, offset, scale, 0.00001)

    params = collections.OrderedDict(zip(params_keys, params_vars))
    moments = collections.OrderedDict(zip(moments_keys, moments_vars))

    return output, params, moments
def bn ( x , params = None , moments = None , backprop_through_moments = True , use_ema = False , is_training = True , ema_epsilon = .9 ) : params_keys , params_vars , moments_keys , moments_vars = [ ] , [ ] , [ ] , [ ] with tf . variable_scope ( 'batch_norm' ) : scope_name = tf . get_variable_scope ( ) . name if use_ema : ema_shape = [ 1 , 1 , 1 , x . get_shape ( ) . as_list ( ) [ - 1 ] ] mean_ema = tf . get_variable ( 'mean_ema' , shape = ema_shape , initializer = tf . initializers . zeros ( ) , trainable = False ) var_ema = tf . get_variable ( 'var_ema' , shape = ema_shape , initializer = tf . initializers . ones ( ) , trainable = False ) if moments is not None : if backprop_through_moments : mean = moments [ scope_name + '/mean' ] var = moments [ scope_name + '/var' ] else : # This variant does not yield good resutls. mean = tf . stop_gradient ( moments [ scope_name + '/mean' ] ) var = tf . stop_gradient ( moments [ scope_name + '/var' ] ) elif use_ema and not is_training : mean = mean_ema var = var_ema else : # If not provided, compute the mean and var of the current batch. replica_ctx = tf . distribute . get_replica_context ( ) if replica_ctx : # from tensorflow/python/keras/layers/normalization/batch_normalization.py  # pylint: disable=line-too-long axes = list ( range ( len ( x . shape ) - 1 ) ) local_sum = tf . reduce_sum ( x , axis = axes , keepdims = True ) local_squared_sum = tf . reduce_sum ( tf . square ( x ) , axis = axes , keepdims = True ) batch_size = tf . cast ( tf . shape ( x ) [ 0 ] , tf . float32 ) x_sum , x_squared_sum , global_batch_size = ( replica_ctx . all_reduce ( 'sum' , [ local_sum , local_squared_sum , batch_size ] ) ) axes_vals = [ ( tf . shape ( x ) ) [ i ] for i in range ( 1 , len ( axes ) ) ] multiplier = tf . cast ( tf . reduce_prod ( axes_vals ) , tf . float32 ) multiplier = multiplier * global_batch_size mean = x_sum / multiplier x_squared_mean = x_squared_sum / multiplier # var = E(x^2) - E(x)^2 var = x_squared_mean - tf . square ( mean ) else : mean , var = tf . nn . moments ( x , axes = list ( range ( len ( x . shape ) - 1 ) ) , keep_dims = True ) # Only update ema's if training and we computed the moments in the current # call.  Note: at test time for episodic learners, ema's may be passed # from the support set to the query set, even if it's not really needed. if use_ema and is_training and moments is None : replica_ctx = tf . distribute . get_replica_context ( ) mean_upd = tf . assign ( mean_ema , mean_ema * ema_epsilon + mean * ( 1.0 - ema_epsilon ) ) var_upd = tf . assign ( var_ema , var_ema * ema_epsilon + var * ( 1.0 - ema_epsilon ) ) updates = tf . group ( [ mean_upd , var_upd ] ) if replica_ctx : tf . add_to_collection ( tf . GraphKeys . UPDATE_OPS , tf . cond ( tf . equal ( replica_ctx . replica_id_in_sync_group , 0 ) , lambda : updates , tf . no_op ) ) else : tf . add_to_collection ( tf . GraphKeys . UPDATE_OPS , updates ) moments_keys += [ scope_name + '/mean' ] moments_vars += [ mean ] moments_keys += [ scope_name + '/var' ] moments_vars += [ var ] if params is None : offset = tf . get_variable ( 'offset' , shape = mean . get_shape ( ) . as_list ( ) , initializer = tf . initializers . zeros ( ) ) scale = tf . get_variable ( 'scale' , shape = var . get_shape ( ) . as_list ( ) , initializer = tf . initializers . ones ( ) ) else : offset = params [ scope_name + '/offset' ] scale = params [ scope_name + '/scale' ] params_keys += [ scope_name + '/offset' ] params_vars += [ offset ] params_keys += [ scope_name + '/scale' ] params_vars += [ scale ] output = tf . nn . batch_normalization ( x , mean , var , offset , scale , 0.00001 ) params = collections . OrderedDict ( zip ( params_keys , params_vars ) ) moments = collections . OrderedDict ( zip ( moments_keys , moments_vars ) ) return output , params , moments