tensorflow_probability/examples/vq_vae.py
def make_encoder(base_depth, activation, latent_size, code_size):
  """Creates the encoder function.

  Args:
    base_depth: Layer base depth in encoder net.
    activation: Activation function in hidden layers.
    latent_size: The number of latent variables in the code.
    code_size: The dimensionality of each latent variable.

  Returns:
    encoder: A `callable` mapping a `Tensor` of images to a `Tensor` of shape
      `[..., latent_size, code_size]`.
  """
  conv = functools.partial(
      tf.keras.layers.Conv2D, padding="SAME", activation=activation)

  encoder_net = tf.keras.Sequential([
      conv(base_depth, 5, 1),
      conv(base_depth, 5, 2),
      conv(2 * base_depth, 5, 1),
      conv(2 * base_depth, 5, 2),
      conv(4 * latent_size, 7, padding="VALID"),
      tf.keras.layers.Flatten(),
      tf.keras.layers.Dense(latent_size * code_size, activation=None),
      tf.keras.layers.Reshape([latent_size, code_size])
  ])

  def encoder(images):
    """Encodes a batch of images.

    Args:
      images: A `Tensor` representing the inputs to be encoded, of shape `[...,
        channels]`.

    Returns:
      codes: A `float`-like `Tensor` of shape `[..., latent_size, code_size]`.
        It represents latent vectors to be matched with the codebook.
    """
    images = 2 * tf.cast(images, dtype=tf.float32) - 1
    codes = encoder_net(images)
    return codes

  return encoder
def make_encoder(base_depth, activation, latent_size, code_size):
  """Creates the encoder function.

  Args:
    base_depth: Layer base depth in encoder net.
    activation: Activation function in hidden layers.
    latent_size: The number of latent variables in the code.
    code_size: The dimensionality of each latent variable.

  Returns:
    encoder: A `callable` mapping a `Tensor` of images to a `Tensor` of shape
      `[..., latent_size, code_size]`.
  """
  conv = functools.partial(
      tf.keras.layers.Conv2D, padding="SAME", activation=activation)

  encoder_net = tf.keras.Sequential([
      conv(base_depth, 5, 1),
      conv(base_depth, 5, 2),
      conv(2 * base_depth, 5, 1),
      conv(2 * base_depth, 5, 2),
      conv(4 * latent_size, 7, padding="VALID"),
      tf.keras.layers.Flatten(),
      tf.keras.layers.Dense(latent_size * code_size, activation=None),
      tf.keras.layers.Reshape([latent_size, code_size])
  ])

  def encoder(images):
    """Encodes a batch of images.

    Args:
      images: A `Tensor` representing the inputs to be encoded, of shape `[...,
        channels]`.

    Returns:
      codes: A `float`-like `Tensor` of shape `[..., latent_size, code_size]`.
        It represents latent vectors to be matched with the codebook.
    """
    images = 2 * tf.cast(images, dtype=tf.float32) - 1
    codes = encoder_net(images)
    return codes

  return encoder
Creates the encoder function.

  Args:
    base_depth: Layer base depth in encoder net.
    activation: Activation function in hidden layers.
    latent_size: The number of latent variables in the code.
    code_size: The dimensionality of each latent variable.

  Returns:
    encoder: A `callable` mapping a `Tensor` of images to a `Tensor` of shape
      `[..., latent_size, code_size]`.