TwoStageVAESampler

Implementation of a Two Stage VAE sampler as proposed in (https://openreview.net/pdf?id=B1e0X3C9tQ).

Available models:

VAE

Vanilla Variational Autoencoder model.

BetaVAE

\(\beta\)-VAE model.

VAE_IAF

Variational Auto Encoder with Inverse Autoregressive Flows (IAF).

DisentangledBetaVAE

Disentangled \(\beta\)-VAE model.

FactorVAE

FactorVAE model.

BetaTCVAE

\(\beta\)-TCVAE model.

IWAE

Importance Weighted Autoencoder model.

MSSSIM_VAE

VAE using perseptual similarity metrics model.

WAE_MMD

Wasserstein Autoencoder model.

INFOVAE_MMD

Info Variational Autoencoder model.

VAMP

Variational Mixture of Posteriors (VAMP) VAE model

VAEGAN

Variational Autoencoder using Adversarial reconstruction loss model.

HVAE

Hamiltonian VAE.

RHVAE

Riemannian Hamiltonian VAE model.
class pythae.samplers.TwoStageVAESamplerConfig(reconstruction_loss='mse', second_stage_depth=2, second_layers_dim=1024)[source]

Two Stage VAE sampler config class.

Parameters

  • latent_dim (int) – The latent space dimension. Default: None.

  • reconstruction_loss (str) – The reconstruction loss to use [‘bce’, ‘mse’]. Default: ‘mse’

  • second_stage_depth (int) – The number of layers in the second stage VAE. Default: 2

  • second_layers_dim (int) – The size of the fully connected layer to used un the second VAE architecture.

class pythae.samplers.TwoStageVAESampler(model, sampler_config=None)[source]

Fits a second VAE in the Autoencoder’s latent space.

Parameters

  • model (VAE) – The VAE model to sample from

  • sampler_config (TwoStageVAESamplerConfig) – A TwoStageVAESamplerConfig instance containing the main parameters of the sampler. If None, a pre-defined configuration is used. Default: None

Note

The method fit must be called to fit the sampler before sampling.

fit(train_data, eval_data=None, training_config=None, batch_size=64)[source]

Method to fit the sampler from the training data

Parameters

  • train_data (Union[torch.Tensor, np.ndarray, Dataset]) – The train data needed to retrieve the training embeddings and fit the second VAE in the latent space.

  • eval_data (Union[torch.Tensor, np.ndarray, Dataset]) – The train data needed to retrieve the evaluation embeddings and fit the second VAE in the latent space.

  • training_config (BaseTrainerConfig) – the training config to use to fit the second VAE.

  • batch_size (int) – The batch size to use to retrieve the embeddings. Default: 64.

sample(num_samples=1, batch_size=500, output_dir=None, return_gen=True, save_sampler_config=False)[source]

Main sampling function of the sampler.

Parameters

  • num_samples (int) – The number of samples to generate

  • batch_size (int) – The batch size to use during sampling

  • output_dir (str) – The directory where the images will be saved. If does not exist the folder is created. If None: the images are not saved. Defaults: None.

  • return_gen (bool) – Whether the sampler should directly return a tensor of generated data. Default: True.

  • save_sampler_config (bool) – Whether to save the sampler config. It is saved in output_dir

Returns

The generated images

Return type

Tensor