Models

This is the heart of Pythae. In here you will find implementations of Autoencoder-based models along with some Normalizing Flows used for improving Variational Inference in the VAE or sampling and Neural Nets to perform benchmark comparison.

Available Autoencoders

BaseAE	Base class for Autoencoder based models.
AutoModel	Utils class allowing to reload any pythae.models automatically
AE	Vanilla Autoencoder model.
VAE	Vanilla Variational Autoencoder model.
BetaVAE	\(\beta\)-VAE model.
VAE_LinNF	Variational Auto Encoder with linear Normalizing Flows 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.
CIWAE	Combination Importance Weighted Autoencoder model.
MIWAE	Multiply Importance Weighted Autoencoder model.
PIWAE	Partially 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
SVAE	\(\mathcal{S}\)-VAE model.
PoincareVAE	Poincaré Variational Autoencoder model.
Adversarial_AE	Adversarial Autoencoder model.
VAEGAN	Variational Autoencoder using Adversarial reconstruction loss model.
VQVAE	Vector Quantized-VAE model.
HVAE	Hamiltonian VAE.
RAE_GP	Regularized Autoencoder with gradient penalty model.
RAE_L2	Regularized Autoencoder with L2 decoder params regularization model.
RHVAE	Riemannian Hamiltonian VAE model.

Available Normalizing Flows

BaseNF	Base Class from Normalizing flows
PlanarFlow	Planar Flow model.
RadialFlow	Radial Flow model.
MADE	Masked Autoencoder model
MAF	Masked Autoregressive Flow.
IAF	Inverse Autoregressive Flow.
PixelCNN	Pixel CNN model.

Basic Example

To launch a model training, you only need to set up your BaseTrainerConfig, BaseAEConfig build the model and trainer accordingly and then call a TrainingPipeline instance.

>>> from pythae.pipelines import TrainingPipeline
>>> from pythae.models import VAE, VAEConfig
>>> from pythae.trainers import BaseTrainerConfig

>>> # Set up the training configuration
>>> my_training_config = BaseTrainerConfig(
...     output_dir='my_model',
...     num_epochs=50,
...     learning_rate=1e-3,
...     per_device_train_batch_size=200,
...     per_device_eval_batch_size=200,
...     steps_saving=None,
...     optimizer_cls="AdamW",
...     optimizer_params={"weight_decay": 0.05, "betas": (0.91, 0.995)},
...     scheduler_cls="ReduceLROnPlateau",
...     scheduler_params={"patience": 5, "factor": 0.5}
... )
>>> # Set up the model configuration
>>> my_vae_config = model_config = VAEConfig(
...     input_dim=(1, 28, 28),
...     latent_dim=10
... )
>>> # Build the model
>>> my_vae_model = VAE(
...     model_config=my_vae_config
... )
>>> # Build the Pipeline
>>> pipeline = TrainingPipeline(
...     training_config=my_training_config,
...     model=my_vae_model
... )
>>> # Launch the Pipeline
>>> pipeline(
...     train_data=your_train_data, # must be torch.Tensor or np.array
...     eval_data=your_eval_data # must be torch.Tensor or np.array
... )

Distributed Training with Pythae

As of v0.1.0, Pythae now supports distributed training using PyTorch’s DDP. It allows you to train your favorite VAE faster and on larger dataset using multi-node and/or multi-gpu training.

To do so, you can built a python script that will then be launched by a launcher (such as srun on a cluster). The only thing that is needed in the script is to specify some elements relative to the environment (such as the number of nodes/gpus) directly in the training configuration as follows

>>> training_config = BaseTrainerConfig(
...     num_epochs=10,
...     learning_rate=1e-3,
...     per_device_train_batch_size=64,
...     per_device_eval_batch_size=64,
...     dist_backend="nccl", # distributed backend
...     world_size=8 # number of gpus to use (n_nodes x n_gpus_per_node),
...     rank=5 # process/gpu id,
...     local_rank=1 # node id,
...     master_addr="localhost" # master address,
...     master_port="1

See this example script that defines a multi-gpu VQVAE training. Be carefull, the way the environnement (world_size, rank …) may be specific to the cluster and launcher you use.

Benchmark

Below are indicated the training times for a Vector Quantized VAE (VQ-VAE) with Pythae for 100 epochs on MNIST on V100 16GB GPU(s), for 50 epochs on FFHQ (1024x1024 images) and for 20 epochs on ImageNet-1k on V100 32GB GPU(s).

Training times of a VQ-VAE with Pythae

Dataset	Data type (train size)	1 GPU	4 GPUs	2x4 GPUs
MNIST	28x28 images (50k)	221.01s	60.32s	34.50s
FFHQ	1024x1024 RGB images (60k)	19h 1min	5h 6min	2h 37min
ImageNet-1k	128x128 RGB images (~ 1.2M)	6h 25min	1h 41min	51min 26s