Add ResNet models and PixelCNN (#19)

* added tests + raise custom errors on wronf flow names in vae_nf_models * Add "VAE_NF" to models __init__ * Add utils import in vae_nf_models * fixed additionnal utils import * Simplified imports in vae_nf_utils * fixed bad import (Encoder_VAE_Conv) * minor changes * structure idea * Add MADE, MAF and correct typos * add NFWrapper * fix flip * add working MAF * add working IAF * black and isort cleanup * doc clean up and addition of NFs * add VAE_IAF + tests * add working VAE_IAF * change gitignor * work on Linear VAE flow * add VAE with linear normalizing flows * Add VAE_LinNF showcase + black + isort * update docs * add demo + VQVAE sampling showcase * update overview notebook * cleaning * first work * add MNIST resnet and clean up tests * add CIFAR and CELEBA resnets * fix test * add sigmoid in resnet * add ema * add resnets * minor fix * minor typo * clean up + pixelCNN * clean up docs * increase nn test coverage * minor changes * clean up imports * apply black * cehck BetaTCVAE * add nll computation in IWAE * check WAE, INFOVAE MSSSI VAE * check VAMP and SVAE * clean up trainers * Check AAE & VAEGAN * Start checking HVAE * check HVAE * update gitignore * udpte docs * update .gitignore * fix typo in trainer * equalize vq vae decoder and encoder with peers * remove no_grad in predict for HVAE * add stopping flag in while * revert trainer zero_grads * fix coupled adversariel trainer test and typos * add predict tests * add AutoModel and tests * fix some imports * fix issue in wae * fix issue in info_vae * add GenerationPipeline * apply black * fix progress bar in eval and speed up eval steps * minor change * fix device issue * fix device issue * update sampler API * add detach in sampler except * add detach in sampler except * change device allocation aae loss * change device allocation vaegan loss * fix IWAE loss issue * fix tqdm on colab * fix tqdm notebook * update .gitignore * update gitignore * allow non square latent dim in vqvae * remove comments * update notebooks * update gitignore * update notebooks * update AutoModel * Uodate notebooks * Update setup and req * Update README and doc * udpate notebooks and gitignore * update notebooks * remove unused import * fix nl hvae * update notebook * Update notebooks * update notebooks * uopdate notebooks * update notebooks * update setup * isort and black Co-authored-by: louis-j-vincent <[email protected]>
clementchadebec · Jun 14, 2022 · e5fe37c · e5fe37c
1 parent 340bb5b
commit e5fe37c
Show file tree

Hide file tree

Showing 239 changed files with 13,552 additions and 2,393 deletions.
diff --git a/.gitignore b/.gitignore
@@ -31,6 +31,23 @@ examples/scripts/reproducibility/*
 wandb
 examples/notebooks/dummy_output_dir/
 examples/notebooks/models_training/dummy_output_dir/
+logs/
+downloads_data/
+*wandb*
+*.slurm
+configs
+examples/scripts/generation_jz.py
+examples/scripts/classification_jz.py
+examples/scripts/latent_dim_sensi.py
+examples/scripts/reconstruction_jz.py
+examples/scripts/jz_array_training.py
+examples/scripts/interpolation_jz.py
+examples/scripts/kmeans_jz.py
+*.csv
+results.ipynb
+examples/scripts/artifacts/*
+examples/scripts/plots/*
+examples/notebooks/my_model_with_custom_archi/
 
 
 

diff --git a/README.md b/README.md
@@ -96,8 +96,8 @@ Below is the list of the models currently implemented in the library.
 | VAMP prior sampler (VAMPSampler)                   |    VAMP   		  | [link](https://arxiv.org/abs/1705.07120) 	  | [link](https://github.com/jmtomczak/vae_vampprior) |
 | Manifold sampler (RHVAESampler)                     |    RHVAE  		  | [link](https://arxiv.org/abs/2105.00026)      |	[link](https://github.com/clementchadebec/pyraug)|
 | Masked Autoregressive Flow Sampler (MAFSampler) | all models | [link](https://arxiv.org/abs/1705.07057v4)      |	[link](https://github.com/gpapamak/maf) |
-| Inverse Autoregressive Flow Sampler (IAFSampler) | all models | [link](https://arxiv.org/abs/1606.04934) |  [link](https://github.com/openai/iaf)             |                     
-
+| Inverse Autoregressive Flow Sampler (IAFSampler) | all models | [link](https://arxiv.org/abs/1606.04934) |  [link](https://github.com/openai/iaf)             |   
+| PixelCNN (PixelCNNSampler) | VQVAE | [link](https://arxiv.org/abs/1606.05328) |             |                     
 
 ## Launching a model training
 
@@ -150,13 +150,47 @@ See [README.md](https://github.com/clementchadebec/benchmark_VAE/tree/main/examp
 
 ## Launching data generation
 
-To launch the data generation process from a trained model, you only need to build your sampler. For instance, to generate new data with your sampler, run the following.
+### Using the `GeneationPipeline`
+
+The easiest way to launch a data generation from a trained model consists in using the built-in `GenerationPipeline` provided in Pythae. Say you want to generate 100 samples using a `MAFSampler` all you have to do is 1) relaod the trained model, 2) define the sampler's configuration and 3) create and launch the `GenerationPipeline` as follows
+
+```python
+>>> from pythae.models import AutoModel
+>>> from pythae.samplers import MAFSamplerConfig
+>>> from pythae.pipelines import GenerationPipeline
+>>> # 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 Samplers
+
+Alternatively, 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.
 
 ```python
->>> from pythae.models import VAE
+>>> from pythae.models import AutoModel
 >>> from pythae.samplers import NormalSampler
 >>> # Retrieve the trained model
->>> my_trained_vae = VAE.load_from_folder(
+>>> my_trained_vae = AutoModel.load_from_folder(
 ...	'path/to/your/trained/model'
 ... )
 >>> # Define your sampler
@@ -175,10 +209,10 @@ If you set `output_dir` to a specific path, the generated images will be saved a
 The samplers can be used with any model as long as it is suited. For instance, a `GaussianMixtureSampler` instance can be used to generate from any model but a `VAMPSampler` will only be usable with a `VAMP` model. Check [here](#available-samplers) to see which ones apply to your model. Be carefull that some samplers such as the `GaussianMixtureSampler` for instance may need to be fitted by calling the `fit` method before using. Below is an example for the `GaussianMixtureSampler`. 
 
 ```python
->>> from pythae.models import VAE
+>>> from pythae.models import AutoModel
 >>> from pythae.samplers import GaussianMixtureSampler, GaussianMixtureSamplerConfig
 >>> # Retrieve the trained model
->>> my_trained_vae = VAE.load_from_folder(
+>>> my_trained_vae = AutoModel.load_from_folder(
 ...	'path/to/your/trained/model'
 ... )
 >>> # Define your sampler
@@ -200,6 +234,7 @@ The samplers can be used with any model as long as it is suited. For instance, a
 ... )
 ```
 
+
 ## Define you own Autoencoder architecture
 
 Pythae 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:
@@ -264,9 +299,9 @@ You can also find predefined neural network architectures for the most common da
 
 ```python
 >>> for pythae.models.nn.benchmark.mnist import (
-...	Encoder_AE_MNIST, # For AE based model (only return embeddings)
-... 	Encoder_VAE_MNIST, # For VAE based model (return embeddings and log_covariances)
-... 	Decoder_AE_MNIST
+...	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
 ... )
 ```
 Replace *mnist* by cifar or celeba to access to other neural nets.

diff --git a/docs/source/advanced/custom_autoencoder.rst → docs/old/advanced/custom_autoencoder.rst b/docs/source/advanced/custom_autoencoder.rst → docs/old/advanced/custom_autoencoder.rst
diff --git a/docs/source/advanced/setting_configs.rst → docs/old/advanced/setting_configs.rst b/docs/source/advanced/setting_configs.rst → docs/old/advanced/setting_configs.rst
diff --git a/docs/source/advanced_use.rst → docs/old/advanced_use.rst b/docs/source/advanced_use.rst → docs/old/advanced_use.rst
diff --git a/docs/source/background.rst → docs/old/background.rst b/docs/source/background.rst → docs/old/background.rst
diff --git a/docs/source/data/pythae.data.datasets.rst → docs/old/data/pythae.data.datasets.rst b/docs/source/data/pythae.data.datasets.rst → docs/old/data/pythae.data.datasets.rst
diff --git a/docs/source/data/pythae.data.loaders.rst → docs/old/data/pythae.data.loaders.rst b/docs/source/data/pythae.data.loaders.rst → docs/old/data/pythae.data.loaders.rst
diff --git a/...source/data/pythae.data.preprocessors.rst → docs/old/data/pythae.data.preprocessors.rst b/...source/data/pythae.data.preprocessors.rst → docs/old/data/pythae.data.preprocessors.rst
diff --git a/docs/source/data/pythae.data.rst → docs/old/data/pythae.data.rst b/docs/source/data/pythae.data.rst → docs/old/data/pythae.data.rst
diff --git a/docs/source/getting_started.rst → docs/old/getting_started.rst b/docs/source/getting_started.rst → docs/old/getting_started.rst
diff --git a/docs/source/math_behind.rst → docs/old/math_behind.rst b/docs/source/math_behind.rst → docs/old/math_behind.rst
diff --git a/docs/source/pythae.config.rst → docs/old/pythae.config.rst b/docs/source/pythae.config.rst → docs/old/pythae.config.rst
diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -15,7 +15,7 @@
 
 # -- Path setup --------------------------------------------------------------
 
-needs_sphinx = "1.6"
+needs_sphinx = "1.3"
 
 sys.path.insert(0, os.path.abspath("../../src/"))
 
@@ -31,7 +31,7 @@
 # Add any Sphinx extension module names here, as strings. They can be
 # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
 # ones.
-extensions = extensions = [
+extensions = [
     "sphinx.ext.autodoc",
     "sphinx.ext.autosummary",
     "sphinx.ext.doctest",
@@ -41,9 +41,12 @@
     "sphinx.ext.napoleon",
     "sphinx.ext.viewcode",
     "sphinx.ext.autosectionlabel",
-    "sphinxcontrib.bibtex",
+    "sphinxcontrib.bibtex"
 ]
 
+
+suppress_warnings = ['autosectionlabel.*']
+
 bibtex_bibfiles = ["references.bib"]
 
 

diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -18,4 +18,33 @@ we can conduct benchmark analysis and reproducible research!
    models/pythae.models
    samplers/pythae.samplers
    trainers/pythae.trainer
-   pipelines/pythae.pipelines
+   pipelines/pythae.pipelines
+
+Setup
+~~~~~~~~~~~~~
+
+To install the latest version of this library run the following using ``pip``
+
+.. code-block:: bash
+
+   $ pip install git+https://github.com/clementchadebec/benchmark_VAE.git 
+
+or alternatively you can clone the github repo to access to tests, tutorials and scripts.
+
+.. code-block:: bash
+
+   $ git clone https://github.com/clementchadebec/benchmark_VAE.git
+
+and install the library
+
+.. code-block:: bash
+
+   $ cd benchmark_VAE
+   $ pip install -e .
+
+If you clone the pythae's repository you will access to  the following:
+
+- ``docs``: The folder in which the documentation can be retrieved.
+- ``tests``: pythae's unit-testing using pytest.
+- ``examples``: A list of ``ipynb`` tutorials and script describing the main functionalities of pythae.
+- ``src/pythae``: The main library which can be installed with ``pip``.
diff --git a/docs/source/models/autoencoders/models.rst b/docs/source/models/autoencoders/models.rst
@@ -1,6 +1,3 @@
-.. _pythae_models:
-
-
 **********************************
 Autoencoders
 **********************************
@@ -31,7 +28,6 @@ Autoencoders
     rae_gp
     rae_l2
     rhvae
-    :nosignatures:
 
 .. automodule::
     pythae.models

diff --git a/docs/source/models/autoencoders/wae_mmd.rst b/docs/source/models/autoencoders/wae_mmd.rst
diff --git a/docs/source/models/nn/celeba/convnets.rst b/docs/source/models/nn/celeba/convnets.rst
@@ -0,0 +1,18 @@
+********************************
+ConvNets
+********************************
+
+.. autoclass:: pythae.models.nn.benchmarks.celeba.Encoder_Conv_AE_CELEBA
+   :members:
+
+.. autoclass:: pythae.models.nn.benchmarks.celeba.Encoder_Conv_VAE_CELEBA
+   :members:
+
+.. autoclass:: pythae.models.nn.benchmarks.celeba.Encoder_Conv_SVAE_CELEBA
+   :members:
+
+.. autoclass:: pythae.models.nn.benchmarks.celeba.Decoder_Conv_AE_CELEBA
+   :members:
+
+.. autoclass:: pythae.models.nn.benchmarks.celeba.Discriminator_Conv_CELEBA
+   :members:
diff --git a/docs/source/models/nn/celeba/pythae_benchmarks_nn_celeba.rst b/docs/source/models/nn/celeba/pythae_benchmarks_nn_celeba.rst
@@ -0,0 +1,35 @@
+**********************************
+CELEBA
+**********************************
+
+.. automodule::
+   pythae.models.nn.benchmarks.celeba
+
+.. toctree::
+   :maxdepth: 1
+
+   convnets
+   resnets
+
+ConvNets
+~~~~~~~~~~~~~~~
+
+.. autosummary::
+   ~pythae.models.nn.benchmarks.celeba.Encoder_Conv_AE_CELEBA
+   ~pythae.models.nn.benchmarks.celeba.Encoder_Conv_VAE_CELEBA
+   ~pythae.models.nn.benchmarks.celeba.Encoder_Conv_SVAE_CELEBA
+   ~pythae.models.nn.benchmarks.celeba.Decoder_Conv_AE_CELEBA
+   ~pythae.models.nn.benchmarks.celeba.Discriminator_Conv_CELEBA
+   :nosignatures:
+
+ResNets
+~~~~~~~~~~~~~~~
+
+.. autosummary::
+   ~pythae.models.nn.benchmarks.celeba.Encoder_ResNet_AE_CELEBA
+   ~pythae.models.nn.benchmarks.celeba.Encoder_ResNet_VAE_CELEBA
+   ~pythae.models.nn.benchmarks.celeba.Encoder_ResNet_SVAE_CELEBA
+   ~pythae.models.nn.benchmarks.celeba.Encoder_ResNet_VQVAE_CELEBA
+   ~pythae.models.nn.benchmarks.celeba.Decoder_ResNet_AE_CELEBA
+   ~pythae.models.nn.benchmarks.celeba.Decoder_ResNet_VQVAE_CELEBA
+   :nosignatures:
diff --git a/docs/source/models/nn/celeba/resnets.rst b/docs/source/models/nn/celeba/resnets.rst
@@ -0,0 +1,21 @@
+********************************
+ResNets
+********************************
+
+.. autoclass:: pythae.models.nn.benchmarks.celeba.Encoder_ResNet_AE_CELEBA
+   :members:
+
+.. autoclass:: pythae.models.nn.benchmarks.celeba.Encoder_ResNet_VAE_CELEBA
+   :members:
+
+.. autoclass:: pythae.models.nn.benchmarks.celeba.Encoder_ResNet_SVAE_CELEBA
+   :members:
+
+.. autoclass:: pythae.models.nn.benchmarks.celeba.Encoder_ResNet_VQVAE_CELEBA
+   :members:
+
+.. autoclass:: pythae.models.nn.benchmarks.celeba.Decoder_ResNet_AE_CELEBA
+   :members:
+
+.. autoclass:: pythae.models.nn.benchmarks.celeba.Decoder_ResNet_VQVAE_CELEBA
+   :members:
diff --git a/docs/source/models/nn/cifar/convnets.rst b/docs/source/models/nn/cifar/convnets.rst
@@ -0,0 +1,18 @@
+********************************
+ConvNets
+********************************
+
+.. autoclass:: pythae.models.nn.benchmarks.cifar.Encoder_Conv_AE_CIFAR
+   :members:
+
+.. autoclass:: pythae.models.nn.benchmarks.cifar.Encoder_Conv_VAE_CIFAR
+   :members:
+
+.. autoclass:: pythae.models.nn.benchmarks.cifar.Encoder_Conv_SVAE_CIFAR
+   :members:
+
+.. autoclass:: pythae.models.nn.benchmarks.cifar.Decoder_Conv_AE_CIFAR
+   :members:
+
+.. autoclass:: pythae.models.nn.benchmarks.cifar.Discriminator_Conv_CIFAR
+   :members:
diff --git a/docs/source/models/nn/cifar/pythae_benchmarks_nn_cifar.rst b/docs/source/models/nn/cifar/pythae_benchmarks_nn_cifar.rst
@@ -0,0 +1,35 @@
+**********************************
+CIFAR
+**********************************
+
+.. automodule::
+   pythae.models.nn.benchmarks.cifar
+
+.. toctree::
+   :maxdepth: 1
+
+   convnets
+   resnets
+
+ConvNets
+~~~~~~~~~~~~~~~
+
+.. autosummary::
+   ~pythae.models.nn.benchmarks.cifar.Encoder_Conv_AE_CIFAR
+   ~pythae.models.nn.benchmarks.cifar.Encoder_Conv_VAE_CIFAR
+   ~pythae.models.nn.benchmarks.cifar.Encoder_Conv_SVAE_CIFAR
+   ~pythae.models.nn.benchmarks.cifar.Decoder_Conv_AE_CIFAR
+   ~pythae.models.nn.benchmarks.cifar.Discriminator_Conv_CIFAR
+   :nosignatures:
+
+ResNets
+~~~~~~~~~~~~~~~
+
+.. autosummary::
+   ~pythae.models.nn.benchmarks.cifar.Encoder_ResNet_AE_CIFAR
+   ~pythae.models.nn.benchmarks.cifar.Encoder_ResNet_VAE_CIFAR
+   ~pythae.models.nn.benchmarks.cifar.Encoder_ResNet_SVAE_CIFAR
+   ~pythae.models.nn.benchmarks.cifar.Encoder_ResNet_VQVAE_CIFAR
+   ~pythae.models.nn.benchmarks.cifar.Decoder_ResNet_AE_CIFAR
+   ~pythae.models.nn.benchmarks.cifar.Decoder_ResNet_VQVAE_CIFAR
+   :nosignatures:
diff --git a/docs/source/models/nn/cifar/resnets.rst b/docs/source/models/nn/cifar/resnets.rst
@@ -0,0 +1,21 @@
+********************************
+ResNets
+********************************
+
+.. autoclass:: pythae.models.nn.benchmarks.cifar.Encoder_ResNet_AE_CIFAR
+   :members:
+
+.. autoclass:: pythae.models.nn.benchmarks.cifar.Encoder_ResNet_VAE_CIFAR
+   :members:
+
+.. autoclass:: pythae.models.nn.benchmarks.cifar.Encoder_ResNet_SVAE_CIFAR
+   :members:
+
+.. autoclass:: pythae.models.nn.benchmarks.cifar.Encoder_ResNet_VQVAE_CIFAR
+   :members:
+
+.. autoclass:: pythae.models.nn.benchmarks.cifar.Decoder_ResNet_AE_CIFAR
+   :members:
+
+.. autoclass:: pythae.models.nn.benchmarks.cifar.Decoder_ResNet_VQVAE_CIFAR
+   :members:
diff --git a/docs/source/models/nn/mnist/convnets.rst b/docs/source/models/nn/mnist/convnets.rst
@@ -0,0 +1,18 @@
+********************************
+ConvNets
+********************************
+
+.. autoclass:: pythae.models.nn.benchmarks.mnist.Encoder_Conv_AE_MNIST
+   :members:
+
+.. autoclass:: pythae.models.nn.benchmarks.mnist.Encoder_Conv_VAE_MNIST
+   :members:
+
+.. autoclass:: pythae.models.nn.benchmarks.mnist.Encoder_Conv_SVAE_MNIST
+   :members:
+
+.. autoclass:: pythae.models.nn.benchmarks.mnist.Decoder_Conv_AE_MNIST
+   :members:
+
+.. autoclass:: pythae.models.nn.benchmarks.mnist.Discriminator_Conv_MNIST
+   :members: