zea.models.hvae

Hierarchical Variational Auto-Encoder for image generation, posterior sampling and inference tasks. To try this model, simply load one of the available presets:

>>> from zea.models.hvae import HierarchicalVAE

>>> model = HierarchicalVAE.from_preset("hvae")

Important

This is a zea implementation of the model. For the original code, see here.

See also

A tutorial notebook where this model is used: Hierarchical VAEs for ultrasound image generation and inpainting.

Classes

HierarchicalVAE(*args, **kwargs)

Hierarchical Variational Autoencoder (HVAE) model.

class zea.models.hvae.HierarchicalVAE(*args, **kwargs)

Bases: DeepGenerativeModel

Hierarchical Variational Autoencoder (HVAE) model. The network as defined here is a snippet of the complete model at: https://github.com/swpenninga/hvae

The lvh versions are trained on EchoNetLVH at 256x256 resolution with 3 channels. (video-frames as channel dimension) The ur(.) versions denote retraining with a UniformRandom agent with (.)/256 lines.

Unlike the other models, this network is built when the weights are loaded.

Parameters:

• name (str) – Name of the model.

• version (str) – Version of the HVAE model to use. Supported versions are: “lvh”, “lvh_ur24”, “lvh_ur16”, “lvh_ur8”, “lvh_ur4”, “lvh_ge24”, “lvh_ge16”, “lvh_ge8”, “lvh_ge4”.

call(measurements)

Returns a reconstruction of the input, together with the latent samples and KL divergences.

Parameters:

measurements (tensor) – Input measurements of shape [B, 256, 256, 3].

Returns:

Reconstructed output of shape [B, 256, 256, 3], List of latent samples from the decoder, and list of KL divergences at each latent layer.

Return type:

recon (tensor)

custom_load_weights(preset)

Load the pretrained weights of the HVAE model from a preset. First builds the model architecture from args.pkl, then loads the weights into the model.

log_density(measurements, **kwargs)

Calculates the log density (ELBO) of the data under the model.

Parameters:

measurements (tensor) – Input measurements of shape [B, 256, 256, 3].

Returns:

negative ELBO of the input measurements, averaged over the batch.

Return type:

-elbo (tensor)

partial_inference(measurements, num_layers=0.5, n_samples=1, **kwargs)

Performs TopDown inference with the HVAE up until a certain layer, after which it continues in the decoder with multiple prior streams.

Parameters:

• measurements (tensor) – Input measurements of shape [B, 256, 256, 3].

• num_layers (float or int) – If float, fraction of total layers to use from the top. If int, number of layers to use from the top.

• n_samples (int) – Number of posterior samples to generate.

Returns:

Posterior samples of shape [B, n_samples, 256, 256, 3].

Return type:

output (tensor)

posterior_sample(measurements, n_samples=1, **kwargs)

Performs posterior sampling on a batch of measurements. Only does a single encoder pass since it is deterministic, but does n_samples decoder passes to create posterior samples.

Parameters:

• measurements (tensor) – Input measurements of shape [B, 256, 256, 3].

• n_samples (int, optional) – Number of posterior samples to generate. Defaults to 1.

Returns:

Posterior samples of shape [B, n_samples, 256, 256, 3].

Return type:

output (tensor)

sample(n_samples=1, **kwargs)

Samples from the prior distribution.

Parameters:

n_samples (int) – Number of samples to generate.

Returns:

Generated samples of shape (n_samples, 256, 256, 3) in [-1, 1].

Return type:

samples (tensor)