tensorflow_probability/examples/disentangled_vae.py
def visualize_qualitative_analysis(inputs, model, samples=1, batch_size=3,
                                   length=8):
  """Visualizes a qualitative analysis of a given model.

  Args:
    inputs: A tensor of the original inputs, of shape [batch, timesteps,
      h, w, c].
    model: A DisentangledSequentialVAE model.
    samples: Number of samples to draw from the latent distributions.
    batch_size: Number of sequences to generate.
    length: Number of timesteps to generate for each sequence.
  """
  average = lambda dist: tf.reduce_mean(
      input_tensor=dist.mean(), axis=0)  # avg over samples
  with tf.compat.v1.name_scope("val_reconstruction"):
    reconstruct = functools.partial(model.reconstruct, inputs=inputs,
                                    samples=samples)
    visualize_reconstruction(inputs, average(reconstruct()))
    visualize_reconstruction(inputs, average(reconstruct(sample_static=True)),
                             name="static_prior")
    visualize_reconstruction(inputs, average(reconstruct(sample_dynamic=True)),
                             name="dynamic_prior")
    visualize_reconstruction(inputs, average(reconstruct(swap_static=True)),
                             name="swap_static")
    visualize_reconstruction(inputs, average(reconstruct(swap_dynamic=True)),
                             name="swap_dynamic")

  with tf.compat.v1.name_scope("generation"):
    generate = functools.partial(model.generate, batch_size=batch_size,
                                 length=length, samples=samples)
    image_summary(average(generate(fix_static=True)), "fix_static")
    image_summary(average(generate(fix_dynamic=True)), "fix_dynamic")
def visualize_qualitative_analysis(inputs, model, samples=1, batch_size=3,
                                   length=8):
  """Visualizes a qualitative analysis of a given model.

  Args:
    inputs: A tensor of the original inputs, of shape [batch, timesteps,
      h, w, c].
    model: A DisentangledSequentialVAE model.
    samples: Number of samples to draw from the latent distributions.
    batch_size: Number of sequences to generate.
    length: Number of timesteps to generate for each sequence.
  """
  average = lambda dist: tf.reduce_mean(
      input_tensor=dist.mean(), axis=0)  # avg over samples
  with tf.compat.v1.name_scope("val_reconstruction"):
    reconstruct = functools.partial(model.reconstruct, inputs=inputs,
                                    samples=samples)
    visualize_reconstruction(inputs, average(reconstruct()))
    visualize_reconstruction(inputs, average(reconstruct(sample_static=True)),
                             name="static_prior")
    visualize_reconstruction(inputs, average(reconstruct(sample_dynamic=True)),
                             name="dynamic_prior")
    visualize_reconstruction(inputs, average(reconstruct(swap_static=True)),
                             name="swap_static")
    visualize_reconstruction(inputs, average(reconstruct(swap_dynamic=True)),
                             name="swap_dynamic")

  with tf.compat.v1.name_scope("generation"):
    generate = functools.partial(model.generate, batch_size=batch_size,
                                 length=length, samples=samples)
    image_summary(average(generate(fix_static=True)), "fix_static")
    image_summary(average(generate(fix_dynamic=True)), "fix_dynamic")
Visualizes a qualitative analysis of a given model.

  Args:
    inputs: A tensor of the original inputs, of shape [batch, timesteps,
      h, w, c].
    model: A DisentangledSequentialVAE model.
    samples: Number of samples to draw from the latent distributions.
    batch_size: Number of sequences to generate.
    length: Number of timesteps to generate for each sequence.