Package: pythae
DocumentationpythaeThis library implements some of the most common (Variational) Autoencoder models under a unified implementation. In particular, it
provides the possibility to perform benchmark experiments and comparisons by training
the models with the same autoencoding neural network architecture. The featuremake your own autoencoderallows you to train any of these models with your own data and own Encoder and Decoder neural networks. It integrates experiment monitoring tools suchwandb,mlfloworcomet-ml🧪 and allows model sharing and loading from theHuggingFace Hub🤗 in a few lines of code.News📢As of v0.1.0,Pythaenow supports distributed training using PyTorch'sDDP. You can now train your favorite VAE faster and on larger datasets, still with a few lines of code.
See our speed-upbenchmark.Quick access:InstallationImplemented models/Implemented samplersReproducibility statement/Results flavorModel training/Data generation/Custom network architectures/Distributed trainingModel sharing with 🤗 Hub/Experiment tracking withwandb/Experiment tracking withmlflow/Experiment tracking withcomet_mlTutorials/DocumentationContributing 🚀/Issues 🛠️Citing this repositoryInstallationTo install the latest stable release of this library run the following usingpip$pipinstallpythaeTo install the latest github version of this library run the following usingpip$pipinstallgit+https://github.com/clementchadebec/benchmark_VAE.gitor alternatively you can clone the github repo to access to tests, tutorials and scripts.$gitclonehttps://github.com/clementchadebec/benchmark_VAE.gitand install the library$cdbenchmark_VAE
$pipinstall-e.Available ModelsBelow is the list of the models currently implemented in the library.ModelsTraining examplePaperOfficial ImplementationAutoencoder (AE)Variational Autoencoder (VAE)linkBeta Variational Autoencoder (BetaVAE)linkVAE with Linear Normalizing Flows (VAE_LinNF)linkVAE with Inverse Autoregressive Flows (VAE_IAF)linklinkDisentangled Beta Variational Autoencoder (DisentangledBetaVAE)linkDisentangling by Factorising (FactorVAE)linkBeta-TC-VAE (BetaTCVAE)linklinkImportance Weighted Autoencoder (IWAE)linklinkMultiply Importance Weighted Autoencoder (MIWAE)linkPartially Importance Weighted Autoencoder (PIWAE)linkCombination Importance Weighted Autoencoder (CIWAE)linkVAE with perceptual metric similarity (MSSSIM_VAE)linkWasserstein Autoencoder (WAE)linklinkInfo Variational Autoencoder (INFOVAE_MMD)linkVAMP Autoencoder (VAMP)linklinkHyperspherical VAE (SVAE)linklinkPoincaré Disk VAE (PoincareVAE)linklinkAdversarial Autoencoder (Adversarial_AE)linkVariational Autoencoder GAN (VAEGAN) 🥗linklinkVector Quantized VAE (VQVAE)linklinkHamiltonian VAE (HVAE)linklinkRegularized AE with L2 decoder param (RAE_L2)linklinkRegularized AE with gradient penalty (RAE_GP)linklinkRiemannian Hamiltonian VAE (RHVAE)linklinkSeereconstructionandgenerationresults for all aforementionned modelsAvailable SamplersBelow is the list of the models currently implemented in the library.SamplersModelsPaperOfficial ImplementationNormal prior (NormalSampler)all modelslinkGaussian mixture (GaussianMixtureSampler)all modelslinklinkTwo stage VAE sampler (TwoStageVAESampler)all VAE based modelslinklinkUnit sphere uniform sampler (HypersphereUniformSampler)SVAElinklinkPoincaré Disk sampler (PoincareDiskSampler)PoincareVAElinklinkVAMP prior sampler (VAMPSampler)VAMPlinklinkManifold sampler (RHVAESampler)RHVAElinklinkMasked Autoregressive Flow Sampler (MAFSampler)all modelslinklinkInverse Autoregressive Flow Sampler (IAFSampler)all modelslinklinkPixelCNN (PixelCNNSampler)VQVAElinkReproducibilityWe validate the implementations by reproducing some results presented in the original publications when the official code has been released or when enough details about the experimental section of the papers were available. Seereproducibilityfor more details.Launching a model trainingTo launch a model training, you only need to call aTrainingPipelineinstance.>>>frompythae.pipelinesimportTrainingPipeline>>>frompythae.modelsimportVAE,VAEConfig>>>frompythae.trainersimportBaseTrainerConfig>>># 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,...train_dataloader_num_workers=2,...eval_dataloader_num_workers=2,...steps_saving=20,...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, np.array or torch datasets...eval_data=your_eval_data# must be torch.Tensor, np.array or torch datasets...)At the end of training, the best model weights, model configuration and training configuration are stored in afinal_modelfolder available inmy_model/MODEL_NAME_training_YYYY-MM-DD_hh-mm-ss(withmy_modelbeing theoutput_dirargument of theBaseTrainerConfig). If you further set thesteps_savingargument to a certain value, folders namedcheckpoint_epoch_kcontaining the best model weights, optimizer, scheduler, configuration and training configuration at epochkwill also appear inmy_model/MODEL_NAME_training_YYYY-MM-DD_hh-mm-ss.Launching a training on benchmark datasetsWe also provide a training script exampleherethat can be used to train the models on benchmarks datasets (mnist, cifar10, celeba ...). The script can be launched with the following commandlinepythontraining.py--datasetmnist--model_nameae--model_config'configs/ae_config.json'--training_config'configs/base_training_config.json'SeeREADME.mdfor further details on this scriptLaunching data generationUsing theGenerationPipelineThe easiest way to launch a data generation from a trained model consists in using the built-inGenerationPipelineprovided in Pythae. Say you want to generate 100 samples using aMAFSamplerall you have to do is 1) relaod the trained model, 2) define the sampler's configuration and 3) create and launch theGenerationPipelineas follows>>>frompythae.modelsimportAutoModel>>>frompythae.samplersimportMAFSamplerConfig>>>frompythae.pipelinesimportGenerationPipeline>>># Retrieve the trained model>>>my_trained_vae=AutoModel.load_from_folder(...'path/to/your/trained/model'...)>>>my_sampler_config=MAFSamplerConfig(...n_made_blocks=2,...n_hidden_in_made=3,...hidden_size=128...)>>># Build the pipeline>>>pipe=GenerationPipeline(...model=my_trained_vae,...sampler_config=my_sampler_config...)>>># Launch data generation>>>generated_samples=pipe(...num_samples=args.num_samples,...return_gen=True,# If false returns nothing...train_data=train_data,# Needed to fit the sampler...eval_data=eval_data,# Needed to fit the sampler...training_config=BaseTrainerConfig(num_epochs=200)# TrainingConfig to use to fit the sampler...)Using the SamplersAlternatively, you can launch the data generation process from a trained model directly with the sampler. For instance, to generate new data with your sampler, run the following.>>>frompythae.modelsimportAutoModel>>>frompythae.samplersimportNormalSampler>>># Retrieve the trained model>>>my_trained_vae=AutoModel.load_from_folder(...'path/to/your/trained/model'...)>>># Define your sampler>>>my_samper=NormalSampler(...model=my_trained_vae...)>>># Generate samples>>>gen_data=my_samper.sample(...num_samples=50,...batch_size=10,...output_dir=None,...return_gen=True...)If you setoutput_dirto a specific path, the generated images will be saved as.pngfiles named00000000.png,00000001.png...
The samplers can be used with any model as long as it is suited. For instance, aGaussianMixtureSamplerinstance can be used to generate from any model but aVAMPSamplerwill only be usable with aVAMPmodel. Checkhereto see which ones apply to your model. Be carefull that some samplers such as theGaussianMixtureSamplerfor instance may need to be fitted by calling thefitmethod before using. Below is an example for theGaussianMixtureSampler.>>>frompythae.modelsimportAutoModel>>>frompythae.samplersimportGaussianMixtureSampler,GaussianMixtureSamplerConfig>>># Retrieve the trained model>>>my_trained_vae=AutoModel.load_from_folder(...'path/to/your/trained/model'...)>>># Define your sampler...gmm_sampler_config=GaussianMixtureSamplerConfig(...n_components=10...)>>>my_samper=GaussianMixtureSampler(...sampler_config=gmm_sampler_config,...model=my_trained_vae...)>>># fit the sampler>>>gmm_sampler.fit(train_dataset)>>># Generate samples>>>gen_data=my_samper.sample(...num_samples=50,...batch_size=10,...output_dir=None,...return_gen=True...)Define you own Autoencoder architecturePythae provides you the possibility to define your own neural networks within the VAE models. For instance, say you want to train a Wassertstein AE with a specific encoder and decoder, you can do the following:>>>frompythae.models.nnimportBaseEncoder,BaseDecoder>>>frompythae.models.base.base_utilsimportModelOutput>>>classMy_Encoder(BaseEncoder):...def__init__(self,args=None):# Args is a ModelConfig instance...BaseEncoder.__init__(self)...self.layers=my_nn_layers()......defforward(self,x:torch.Tensor)->ModelOutput:...out=self.layers(x)...output=ModelOutput(...embedding=out# Set the output from the encoder in a ModelOutput instance...)...returnoutput......classMy_Decoder(BaseDecoder):...def__init__(self,args=None):...BaseDecoder.__init__(self)...self.layers=my_nn_layers()......defforward(self,x:torch.Tensor)->ModelOutput:...out=self.layers(x)...output=ModelOutput(...reconstruction=out# Set the output from the decoder in a ModelOutput instance...)...returnoutput...>>>my_encoder=My_Encoder()>>>my_decoder=My_Decoder()And now build the model>>>frompythae.modelsimportWAE_MMD,WAE_MMD_Config>>># Set up the model configuration>>>my_wae_config=model_config=WAE_MMD_Config(...input_dim=(1,28,28),...latent_dim=10...)...>>># Build the model>>>my_wae_model=WAE_MMD(...model_config=my_wae_config,...encoder=my_encoder,# pass your encoder as argument when building the model...decoder=my_decoder# pass your decoder as argument when building the model...)important note 1: For all AE-based models (AE, WAE, RAE_L2, RAE_GP), both the encoder and decoder must return aModelOutputinstance. For the encoder, theModelOutputinstance must contain the embbeddings under the keyembedding. For the decoder, theModelOutputinstance must contain the reconstructions under the keyreconstruction.important note 2: For all VAE-based models (VAE, BetaVAE, IWAE, HVAE, VAMP, RHVAE), both the encoder and decoder must return aModelOutputinstance. For the encoder, theModelOutputinstance must contain the embbeddings andlog-covariance matrices (of shape batch_size x latent_space_dim) respectively under the keyembeddingandlog_covariancekey. For the decoder, theModelOutputinstance must contain the reconstructions under the keyreconstruction.Using benchmark neural netsYou can also find predefined neural network architectures for the most common data sets (i.e.MNIST, CIFAR, CELEBA ...) that can be loaded as follows>>>frompythae.models.nn.benchmark.mnistimport(...Encoder_Conv_AE_MNIST,# For AE based model (only return embeddings)...Encoder_Conv_VAE_MNIST,# For VAE based model (return embeddings and log_covariances)...Decoder_Conv_AE_MNIST...)Replacemnistby cifar or celeba to access to other neural nets.Distributed Training withPythaeAs ofv0.1.0, Pythae now supports distributed training using PyTorch'sDDP. It allows you to train your favorite VAE faster and on larger dataset using multi-gpu and/or multi-node training.To do so, you can build a python script that will then be launched by a launcher (such assrunon a cluster). The only thing that is needed in the script is to specify some elements relative to the distributed 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,...train_dataloader_num_workers=8,...eval_dataloader_num_workers=8,...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="12345"# master port,...)See thisexample scriptthat defines a multi-gpu VQVAE training on ImageNet dataset. Please note that the way the distributed environnement variables (world_size,rank...) are recovered may be specific to the cluster and launcher you use.BenchmarkBelow are indicated the training times for a Vector Quantized VAE (VQ-VAE) withPythaefor 100 epochs on MNIST on V100 16GB GPU(s), for 50 epochs onFFHQ(1024x1024 images) and for 20 epochs onImageNet-1kon V100 32GB GPU(s).Train Data1 GPU4 GPUs2x4 GPUsMNIST (VQ-VAE)28x28 images (50k)235.18 s62.00 s35.86 sFFHQ 1024x1024 (VQVAE)1024x1024 RGB images (60k)19h 1min5h 6min2h 37minImageNet-1k 128x128 (VQVAE)128x128 RGB images (~ 1.2M)6h 25min1h 41min51min 26sFor each dataset, we provide the benchmarking scriptshereSharing your models with the HuggingFace Hub 🤗Pythae also allows you to share your models on theHuggingFace Hub. To do so you need:a valid HuggingFace accountthe packagehuggingface_hubinstalled in your virtual env. If not you can install it with$ python -m pip install huggingface_hubto be logged in to your HuggingFace account using$ huggingface-cli loginUploading a model to the HubAny pythae model can be easily uploaded using the methodpush_to_hf_hub>>>my_vae_model.push_to_hf_hub(hf_hub_path="your_hf_username/your_hf_hub_repo")Note:Ifyour_hf_hub_repoalready exists and is not empty, files will be overridden. In case,
the repoyour_hf_hub_repodoes not exist, a folder having the same name will be created.Downloading models from the HubEquivalently, you can download or reload any Pythae's model directly from the Hub using the methodload_from_hf_hub>>>frompythae.modelsimportAutoModel>>>my_downloaded_vae=AutoModel.load_from_hf_hub(hf_hub_path="path_to_hf_repo")Monitoring your experiments withwandb🧪Pythae also integrates the experiment tracking toolwandballowing users to store their configs, monitor their trainings and compare runs through a graphic interface. To be able use this feature you will need:a valid wandb accountthe packagewandbinstalled in your virtual env. If not you can install it with$ pip install wandbto be logged in to your wandb account using$ wandb loginCreating aWandbCallbackLaunching an experiment monitoring withwandbin pythae is pretty simple. The only thing a user needs to do is create aWandbCallbackinstance...>>># Create you callback>>>frompythae.trainers.training_callbacksimportWandbCallback>>>callbacks=[]# the TrainingPipeline expects a list of callbacks>>>wandb_cb=WandbCallback()# Build the callback>>># SetUp the callback>>>wandb_cb.setup(...training_config=your_training_config,# training config...model_config=your_model_config,# model config...project_name="your_wandb_project",# specify your wandb project...entity_name="your_wandb_entity",# specify your wandb entity...)>>>callbacks.append(wandb_cb)# Add it to the callbacks list...and then pass it to theTrainingPipeline.>>>pipeline=TrainingPipeline(...training_config=config,...model=model...)>>>pipeline(...train_data=train_dataset,...eval_data=eval_dataset,...callbacks=callbacks# pass the callbacks to the TrainingPipeline and you are done!...)>>># You can log to https://wandb.ai/your_wandb_entity/your_wandb_project to monitor your trainingSee the detailed tutorialMonitoring your experiments withmlflow🧪Pythae also integrates the experiment tracking toolmlflowallowing users to store their configs, monitor their trainings and compare runs through a graphic interface. To be able use this feature you will need:the packagemlfowinstalled in your virtual env. If not you can install it with$ pip install mlflowCreating aMLFlowCallbackLaunching an experiment monitoring withmlfowin pythae is pretty simple. The only thing a user needs to do is create aMLFlowCallbackinstance...>>># Create you callback>>>frompythae.trainers.training_callbacksimportMLFlowCallback>>>callbacks=[]# the TrainingPipeline expects a list of callbacks>>>mlflow_cb=MLFlowCallback()# Build the callback>>># SetUp the callback>>>mlflow_cb.setup(...training_config=your_training_config,# training config...model_config=your_model_config,# model config...run_name="mlflow_cb_example",# specify your mlflow run...)>>>callbacks.append(mlflow_cb)# Add it to the callbacks list...and then pass it to theTrainingPipeline.>>>pipeline=TrainingPipeline(...training_config=config,...model=model...)>>>pipeline(...train_data=train_dataset,...eval_data=eval_dataset,...callbacks=callbacks# pass the callbacks to the TrainingPipeline and you are done!...)you can visualize your metric by running the following in the directory where the./mlruns$mlflowuiSee the detailed tutorialMonitoring your experiments withcomet_ml🧪Pythae also integrates the experiment tracking toolcomet_mlallowing users to store their configs, monitor their trainings and compare runs through a graphic interface. To be able use this feature you will need:the packagecomet_mlinstalled in your virtual env. If not you can install it with$ pip install comet_mlCreating aCometCallbackLaunching an experiment monitoring withcomet_mlin pythae is pretty simple. The only thing a user needs to do is create aCometCallbackinstance...>>># Create you callback>>>frompythae.trainers.training_callbacksimportCometCallback>>>callbacks=[]# the TrainingPipeline expects a list of callbacks>>>comet_cb=CometCallback()# Build the callback>>># SetUp the callback>>>comet_cb.setup(...training_config=training_config,# training config...model_config=model_config,# model config...api_key="your_comet_api_key",# specify your comet api-key...project_name="your_comet_project",# specify your wandb project...#offline_run=True, # run in offline mode...#offline_directory='my_offline_runs' # set the directory to store the offline runs...)>>>callbacks.append(comet_cb)# Add it to the callbacks list...and then pass it to theTrainingPipeline.>>>pipeline=TrainingPipeline(...training_config=config,...model=model...)>>>pipeline(...train_data=train_dataset,...eval_data=eval_dataset,...callbacks=callbacks# pass the callbacks to the TrainingPipeline and you are done!...)>>># You can log to https://comet.com/your_comet_username/your_comet_project to monitor your trainingSee the detailed tutorialGetting your hands on the codeTo help you to understand the way pythae works and how you can train your models with this library we also
provide tutorials:making_your_own_autoencoder.ipynbshows you how to pass your own networks to the models implemented in pythaecustom_dataset.ipynbshows you how to  use custom datasets with any of the models implemented in pythaehf_hub_models_sharing.ipynbshows you how to upload and download models for the HuggingFace Hubwandb_experiment_monitoring.ipynbshows you how to monitor you experiments usingwandbmlflow_experiment_monitoring.ipynbshows you how to monitor you experiments usingmlflowcomet_experiment_monitoring.ipynbshows you how to monitor you experiments usingcomet_mlmodels_trainingfolder provides notebooks showing how to train each implemented model and how to sample from it usingpythae.samplers.scriptsfolder provides in particular an example of a training script to train the models on benchmark data sets (mnist, cifar10, celeba ...)Dealing with issues 🛠️If you are experiencing any issues while running the code or request new features/models to be implemented pleaseopen an issue on github.Contributing 🚀You want to contribute to this library by adding a model, a sampler or simply fix a bug ? That's awesome! Thank you! Please seeCONTRIBUTING.mdto follow the main contributing guidelines.ResultsReconstructionFirst let's have a look at the reconstructed samples taken from the evaluation set.ModelsMNISTCELEBAEval dataAEVAEBeta-VAEVAE Lin NFVAE IAFDisentangled  Beta-VAEFactorVAEBetaTCVAEIWAEMSSSIM_VAEWAEINFO VAEVAMPSVAEAdversarial_AEVAE_GANVQVAEHVAERAE_L2RAE_GPRiemannian Hamiltonian VAE (RHVAE)GenerationHere, we show the generated samples using each model implemented in the library and different samplers.ModelsMNISTCELEBAAE  + GaussianMixtureSamplerVAE  + NormalSamplerVAE  + GaussianMixtureSamplerVAE  + TwoStageVAESamplerVAE  + MAFSamplerBeta-VAE + NormalSamplerVAE Lin NF + NormalSamplerVAE IAF + NormalSamplerDisentangled Beta-VAE + NormalSamplerFactorVAE + NormalSamplerBetaTCVAE + NormalSamplerIWAE +  Normal samplerMSSSIM_VAE  + NormalSamplerWAE + NormalSamplerINFO VAE + NormalSamplerSVAE + HypershereUniformSamplerVAMP + VAMPSamplerAdversarial_AE + NormalSamplerVAEGAN + NormalSamplerVQVAE + MAFSamplerHVAE + NormalSamplerRAE_L2 + GaussianMixtureSamplerRAE_GP + GaussianMixtureSamplerRiemannian Hamiltonian VAE (RHVAE) + RHVAE SamplerCitationIf you find this work useful or use it in your research, please consider citing us@inproceedings{chadebec2022pythae,author={Chadebec, Cl\'{e}ment and Vincent, Louis and Allassonniere, Stephanie},booktitle={Advances in Neural Information Processing Systems},editor={S. Koyejo and S. Mohamed and A. Agarwal and D. Belgrave and K. Cho and A. Oh},pages={21575--21589},publisher={Curran Associates, Inc.},title={Pythae: Unifying Generative Autoencoders in Python - A Benchmarking Use Case},volume={35},year={2022}}