(x,
                   film_selector,
                   num_sets,
                   film_weight_decay,
                   params=None,
                   moments=None,
                   is_training=True,
                   backprop_through_moments=True)
Batch normalization used during training FLUTE's template.
Batch normalization used during training FLUTE's template.
Batch normalization used during training FLUTE s template .
def bn_flute_train(x,
                   film_selector,
                   num_sets,
                   film_weight_decay,
                   params=None,
                   moments=None,
                   is_training=True,
                   backprop_through_moments=True):
  """Batch normalization used during training FLUTE's template."""
  del is_training, backprop_through_moments  # Not needed.
  params_keys, params_vars, moments_keys, moments_vars = [], [], [], []
  with tf.variable_scope('batch_norm'):
    scope_name = tf.get_variable_scope().name

    # Part 1: Get the statistics (mean and var) to use during normalization.
    # Compute the mean and var of the current batch. [1, 1, 1, num channels].
    mean, var = tf.nn.moments(
        x, axes=list(range(len(x.shape) - 1)), keep_dims=True)
    num_channels = mean.shape[-1]

    if moments is not None:
      # A common use case for this: passing in the moments computed from the
      # support set during the query set forward pass.
      mean = moments[scope_name + '/mean']
      var = moments[scope_name + '/var']

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

    # Part 2: Select the scale and offset params to use during normalization.
    if params is None:
      offset = tf.get_variable(
          'offset',
          shape=[num_sets, num_channels],
          initializer=tf.initializers.zeros(),
          regularizer=tf.keras.regularizers.L2(film_weight_decay))
      # We init the scale to zeros and then add 1 (instead of init at 1) so that
      # the weight decay regularizes the scale parameters towards unity.
      scale = tf.get_variable(
          'scale',
          shape=[num_sets, num_channels],
          initializer=tf.initializers.zeros(),
          regularizer=tf.keras.regularizers.L2(film_weight_decay))
      scale = scale + 1

      # Combine the parameters according to the given selector.
      # pylint: disable=g-explicit-length-test
      if isinstance(film_selector, int) or not len(film_selector.shape):
        # pylint: enable=g-explicit-length-test
        # If sel is an int or []-shaped int Tensor, expand it to [1, num_sets].
        film_selector = tf.one_hot(tf.cast(film_selector, tf.int64), num_sets)
        film_selector = tf.expand_dims(film_selector, 0)
        film_selector = tf.cast(film_selector, tf.float32)
      # [1, num_sets] x [num_sets, num channels] -> [1, num channels].
      scale = tf.matmul(film_selector, scale)
      offset = tf.matmul(film_selector, offset)
      # Reshape to [1, 1, 1, num_channels].
      scale = tf.reshape(scale, [1, 1, 1, -1])
      offset = tf.reshape(offset, [1, 1, 1, -1])

    else:
      offset = tf.squeeze(params[scope_name + '/offset'])
      scale = tf.squeeze(params[scope_name + '/scale'])

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

    # Part 3: Perform batch normalization with the selected stats and params.
    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_flute_train ( x , film_selector , num_sets , film_weight_decay , params = None , moments = None , is_training = True , backprop_through_moments = True ) : del is_training , backprop_through_moments # Not needed. params_keys , params_vars , moments_keys , moments_vars = [ ] , [ ] , [ ] , [ ] with tf . variable_scope ( 'batch_norm' ) : scope_name = tf . get_variable_scope ( ) . name # Part 1: Get the statistics (mean and var) to use during normalization. # Compute the mean and var of the current batch. [1, 1, 1, num channels]. mean , var = tf . nn . moments ( x , axes = list ( range ( len ( x . shape ) - 1 ) ) , keep_dims = True ) num_channels = mean . shape [ - 1 ] if moments is not None : # A common use case for this: passing in the moments computed from the # support set during the query set forward pass. mean = moments [ scope_name + '/mean' ] var = moments [ scope_name + '/var' ] moments_keys += [ scope_name + '/mean' ] moments_vars += [ mean ] moments_keys += [ scope_name + '/var' ] moments_vars += [ var ] # Part 2: Select the scale and offset params to use during normalization. if params is None : offset = tf . get_variable ( 'offset' , shape = [ num_sets , num_channels ] , initializer = tf . initializers . zeros ( ) , regularizer = tf . keras . regularizers . L2 ( film_weight_decay ) ) # We init the scale to zeros and then add 1 (instead of init at 1) so that # the weight decay regularizes the scale parameters towards unity. scale = tf . get_variable ( 'scale' , shape = [ num_sets , num_channels ] , initializer = tf . initializers . zeros ( ) , regularizer = tf . keras . regularizers . L2 ( film_weight_decay ) ) scale = scale + 1 # Combine the parameters according to the given selector. # pylint: disable=g-explicit-length-test if isinstance ( film_selector , int ) or not len ( film_selector . shape ) : # pylint: enable=g-explicit-length-test # If sel is an int or []-shaped int Tensor, expand it to [1, num_sets]. film_selector = tf . one_hot ( tf . cast ( film_selector , tf . int64 ) , num_sets ) film_selector = tf . expand_dims ( film_selector , 0 ) film_selector = tf . cast ( film_selector , tf . float32 ) # [1, num_sets] x [num_sets, num channels] -> [1, num channels]. scale = tf . matmul ( film_selector , scale ) offset = tf . matmul ( film_selector , offset ) # Reshape to [1, 1, 1, num_channels]. scale = tf . reshape ( scale , [ 1 , 1 , 1 , - 1 ] ) offset = tf . reshape ( offset , [ 1 , 1 , 1 , - 1 ] ) else : offset = tf . squeeze ( params [ scope_name + '/offset' ] ) scale = tf . squeeze ( params [ scope_name + '/scale' ] ) params_keys += [ scope_name + '/offset' ] params_vars += [ offset ] params_keys += [ scope_name + '/scale' ] params_vars += [ scale ] # Part 3: Perform batch normalization with the selected stats and params. 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