DisentangledBetaVAE

This module is the implementation of the Disentangled Beta VAE proposed in (https://arxiv.org/abs/1804.03599). This model adds a new parameter to the \(\beta\)-VAE loss function corresponding to the target value for the KL between the prior and the posterior distribution. It is progressively increased throughout training.

Available samplers

NormalSampler

Samples from a Standard normal distribution in the Autoencoder’s latent space.

GaussianMixtureSampler

Fits a Gaussian Mixture in the Autoencoder’s latent space.

TwoStageVAESampler

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

MAFSampler

Fits a Masked Autoregressive Flow in the Autoencoder’s latent space.

IAFSampler

Fits an Inverse Autoregressive Flow in the Autoencoder’s latent space.

class pythae.models.DisentangledBetaVAEConfig(input_dim=None, latent_dim=10, uses_default_encoder=True, uses_default_decoder=True, reconstruction_loss='mse', beta=10.0, C=50.0, warmup_epoch=25)[source]

Disentangled \(\beta\)-VAE model config config class

Parameters
  • input_dim (tuple) – The input_data dimension.

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

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

  • beta (float) – The balancing factor. Default: 10.

  • C (float) – The value of the KL divergence term of the ELBO we wish to approach, measured in nats. Default: 50.

  • warmup_epoch (int) – The number of epochs during which the KL divergence objective will increase from 0 to C (should be smaller or equal to nb_epochs). Default: 100

  • epoch (int) – The current epoch. Default: 0

class pythae.models.DisentangledBetaVAE(model_config, encoder=None, decoder=None)[source]

Disentangled \(\beta\)-VAE model.

Parameters
  • model_config (DisentangledBetaVAEConfig) – The Variational Autoencoder configuration setting the main parameters of the model.

  • encoder (BaseEncoder) – An instance of BaseEncoder (inheriting from torch.nn.Module which plays the role of encoder. This argument allows you to use your own neural networks architectures if desired. If None is provided, a simple Multi Layer Preception (https://en.wikipedia.org/wiki/Multilayer_perceptron) is used. Default: None.

  • decoder (BaseDecoder) – An instance of BaseDecoder (inheriting from torch.nn.Module which plays the role of decoder. This argument allows you to use your own neural networks architectures if desired. If None is provided, a simple Multi Layer Preception (https://en.wikipedia.org/wiki/Multilayer_perceptron) is used. Default: None.

Note

For high dimensional data we advice you to provide you own network architectures. With the provided MLP you may end up with a MemoryError.

forward(inputs, **kwargs)[source]

The VAE model

Parameters

inputs (BaseDataset) – The training dataset with labels

Returns

An instance of ModelOutput containing all the relevant parameters

Return type

ModelOutput