tensorflow_probability/examples/vq_vae.py
def add_ema_control_dependencies(vector_quantizer,
                                 one_hot_assignments,
                                 codes,
                                 commitment_loss,
                                 decay):
  """Add control dependencies to the commmitment loss to update the codebook.

  Args:
    vector_quantizer: An instance of the VectorQuantizer class.
    one_hot_assignments: The one-hot vectors corresponding to the matched
      codebook entry for each code in the batch.
    codes: A `float`-like `Tensor` containing the latent vectors to be compared
      to the codebook.
    commitment_loss: The commitment loss from comparing the encoder outputs to
      their neighboring codebook entries.
    decay: Decay factor for exponential moving average.

  Returns:
    commitment_loss: Commitment loss with control dependencies.
  """
  # Use an exponential moving average to update the codebook.
  updated_ema_count = moving_averages.assign_moving_average(
      vector_quantizer.ema_count,
      tf.reduce_sum(input_tensor=one_hot_assignments, axis=[0, 1]),
      decay,
      zero_debias=False)
  updated_ema_means = moving_averages.assign_moving_average(
      vector_quantizer.ema_means,
      tf.reduce_sum(
          input_tensor=tf.expand_dims(codes, 2) *
          tf.expand_dims(one_hot_assignments, 3),
          axis=[0, 1]),
      decay,
      zero_debias=False)

  # Add small value to avoid dividing by zero.
  perturbed_ema_count = updated_ema_count + 1e-5
  with tf.control_dependencies([commitment_loss]):
    update_means = tf.compat.v1.assign(
        vector_quantizer.codebook,
        updated_ema_means / perturbed_ema_count[..., tf.newaxis])
    with tf.control_dependencies([update_means]):
      return tf.identity(commitment_loss)
def add_ema_control_dependencies(vector_quantizer,
                                 one_hot_assignments,
                                 codes,
                                 commitment_loss,
                                 decay):
  # Use an exponential moving average to update the codebook.
  updated_ema_count = moving_averages.assign_moving_average(
      vector_quantizer.ema_count,
      tf.reduce_sum(input_tensor=one_hot_assignments, axis=[0, 1]),
      decay,
      zero_debias=False)
  updated_ema_means = moving_averages.assign_moving_average(
      vector_quantizer.ema_means,
      tf.reduce_sum(
          input_tensor=tf.expand_dims(codes, 2) *
          tf.expand_dims(one_hot_assignments, 3),
          axis=[0, 1]),
      decay,
      zero_debias=False)

  # Add small value to avoid dividing by zero.
  perturbed_ema_count = updated_ema_count + 1e-5
  with tf.control_dependencies([commitment_loss]):
    update_means = tf.compat.v1.assign(
        vector_quantizer.codebook,
        updated_ema_means / perturbed_ema_count[..., tf.newaxis])
    with tf.control_dependencies([update_means]):
      return tf.identity(commitment_loss)
def add_ema_control_dependencies ( vector_quantizer , one_hot_assignments , codes , commitment_loss , decay ) : # Use an exponential moving average to update the codebook. updated_ema_count = moving_averages . assign_moving_average ( vector_quantizer . ema_count , tf . reduce_sum ( input_tensor = one_hot_assignments , axis = [ 0 , 1 ] ) , decay , zero_debias = False ) updated_ema_means = moving_averages . assign_moving_average ( vector_quantizer . ema_means , tf . reduce_sum ( input_tensor = tf . expand_dims ( codes , 2 ) * tf . expand_dims ( one_hot_assignments , 3 ) , axis = [ 0 , 1 ] ) , decay , zero_debias = False ) # Add small value to avoid dividing by zero. perturbed_ema_count = updated_ema_count + 1e-5 with tf . control_dependencies ( [ commitment_loss ] ) : update_means = tf . compat . v1 . assign ( vector_quantizer . codebook , updated_ema_means / perturbed_ema_count [ ... , tf . newaxis ] ) with tf . control_dependencies ( [ update_means ] ) : return tf . identity ( commitment_loss )
Add control dependencies to the commmitment loss to update the codebook.

  Args:
    vector_quantizer: An instance of the VectorQuantizer class.
    one_hot_assignments: The one-hot vectors corresponding to the matched
      codebook entry for each code in the batch.
    codes: A `float`-like `Tensor` containing the latent vectors to be compared
      to the codebook.
    commitment_loss: The commitment loss from comparing the encoder outputs to
      their neighboring codebook entries.
    decay: Decay factor for exponential moving average.

  Returns:
    commitment_loss: Commitment loss with control dependencies.
Add control dependencies to the commmitment loss to update the codebook .