(x,
                  film_selector,
                  num_sets,
                  film_weight_decay=None,
                  params=None,
                  moments=None,
                  is_training=True,
                  backprop_through_moments=True)
Batch normalization used for FLUTE evaluation.
Batch normalization used for FLUTE evaluation.
Batch normalization used for FLUTE evaluation .
def bn_flute_eval(x,
                  film_selector,
                  num_sets,
                  film_weight_decay=None,
                  params=None,
                  moments=None,
                  is_training=True,
                  backprop_through_moments=True):
  """Batch normalization used for FLUTE evaluation."""
  del is_training, backprop_through_moments, film_weight_decay  # Not needed.
  params_keys, params_vars, moments_keys, moments_vars = [], [], [], []

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

    if moments is None:
      # 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)
    else:
      mean = moments[scope_name + '/mean']
      var = moments[scope_name + '/var']

    num_channels = mean.shape[-1]
    moments_keys += [scope_name + '/mean']
    moments_vars += [mean]
    moments_keys += [scope_name + '/var']
    moments_vars += [var]

    if params is None:
      # These variables will be restored from the pre-trained checkpoint.
      offset = tf.get_variable(
          'offset',
          shape=[num_sets, num_channels],
          initializer=tf.initializers.zeros())
      scale = tf.get_variable(
          'scale',
          shape=[num_sets, num_channels],
          initializer=tf.initializers.zeros())
      scale = scale + 1

      if film_selector is not None:
        # 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 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])

      # The following two variables are the ones that FiLMLearner will optimize
      # for each task using the support set. These are additive offsets to the
      # selected initialization (e.g. the proposed blending of the per-dataset
      # sets of FiLM parameters).
      # [1, 1, 1, num_channels].
      offset_for_film_learner = tf.get_variable(
          'offset_for_film_learner',
          shape=mean.get_shape().as_list(),
          initializer=tf.initializers.zeros())
      scale_for_film_learner = tf.get_variable(
          'scale_for_film_learner',
          shape=var.get_shape().as_list(),
          initializer=tf.initializers.zeros())

      if film_selector is None:
        # Learn a set of FiLM parameters from scratch, ignoring reloaded ones.
        offset = offset_for_film_learner
        scale = scale_for_film_learner + 1
      else:
        # The new offset/scale are additive offsets to the reloaded ones.
        offset = offset_for_film_learner + offset
        scale = scale_for_film_learner + scale

    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_flute_eval ( x , film_selector , num_sets , film_weight_decay = None , params = None , moments = None , is_training = True , backprop_through_moments = True ) : del is_training , backprop_through_moments , film_weight_decay # Not needed. params_keys , params_vars , moments_keys , moments_vars = [ ] , [ ] , [ ] , [ ] with tf . variable_scope ( 'batch_norm' ) : scope_name = tf . get_variable_scope ( ) . name if moments is None : # 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 ) else : mean = moments [ scope_name + '/mean' ] var = moments [ scope_name + '/var' ] num_channels = mean . shape [ - 1 ] moments_keys += [ scope_name + '/mean' ] moments_vars += [ mean ] moments_keys += [ scope_name + '/var' ] moments_vars += [ var ] if params is None : # These variables will be restored from the pre-trained checkpoint. offset = tf . get_variable ( 'offset' , shape = [ num_sets , num_channels ] , initializer = tf . initializers . zeros ( ) ) scale = tf . get_variable ( 'scale' , shape = [ num_sets , num_channels ] , initializer = tf . initializers . zeros ( ) ) scale = scale + 1 if film_selector is not None : # 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 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 ] ) # The following two variables are the ones that FiLMLearner will optimize # for each task using the support set. These are additive offsets to the # selected initialization (e.g. the proposed blending of the per-dataset # sets of FiLM parameters). # [1, 1, 1, num_channels]. offset_for_film_learner = tf . get_variable ( 'offset_for_film_learner' , shape = mean . get_shape ( ) . as_list ( ) , initializer = tf . initializers . zeros ( ) ) scale_for_film_learner = tf . get_variable ( 'scale_for_film_learner' , shape = var . get_shape ( ) . as_list ( ) , initializer = tf . initializers . zeros ( ) ) if film_selector is None : # Learn a set of FiLM parameters from scratch, ignoring reloaded ones. offset = offset_for_film_learner scale = scale_for_film_learner + 1 else : # The new offset/scale are additive offsets to the reloaded ones. offset = offset_for_film_learner + offset scale = scale_for_film_learner + scale 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