Multiply/Partially/Combination IWAE (#54)

* work on CIWAE and MIWAE * add PIWAE * minor change in Coupled Optimizer Trainer * update repro script * prepare repro * add DynBinDataset to scripts * fix number of training samples * update test calbacks with trainer update * fix issue in callback * add eval model * add eval mode * add examples and update readme * add reproduction results * apply black and isort * clean up scripts * clean up scripts * clean up scripts * removed unecessary * prepare release
clementchadebec · Sep 7, 2022 · 2a782b9 · 2a782b9
1 parent 454eae9
commit 2a782b9
Show file tree

Hide file tree

Showing 54 changed files with 5,220 additions and 68 deletions.
diff --git a/.gitignore b/.gitignore
@@ -59,6 +59,8 @@ examples/scripts/configs/*
 examples/notebooks/mlruns/*
 mlruns/*
 examples/scripts/reproducibility/reproducibility/*
+examples/scripts/reproducibility/sync.ipynb
+
 
 
 

diff --git a/README.md b/README.md
@@ -86,7 +86,10 @@ VAE with Inverse Autoregressive Flows (VAE_IAF) |  [![Open In Colab](https://col
 | Disentangled Beta Variational Autoencoder (DisentangledBetaVAE) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/clementchadebec/benchmark_VAE/blob/main/examples/notebooks/models_training/disentangled_beta_vae_training.ipynb) | [link](https://arxiv.org/abs/1804.03599)  |   
 | Disentangling by Factorising (FactorVAE) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/clementchadebec/benchmark_VAE/blob/main/examples/notebooks/models_training/factor_vae_training.ipynb) | [link](https://arxiv.org/abs/1802.05983)  |                                                                            |
 | Beta-TC-VAE (BetaTCVAE) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/clementchadebec/benchmark_VAE/blob/main/examples/notebooks/models_training/beta_tc_vae_training.ipynb) | [link](https://arxiv.org/abs/1802.04942)  |  [link](https://github.com/rtqichen/beta-tcvae)
-| Importance Weighted Autoencoder (IWAE) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/clementchadebec/benchmark_VAE/blob/main/examples/notebooks/models_training/iwae_training.ipynb) | [link](https://arxiv.org/abs/1509.00519v4)  | [link](https://github.com/yburda/iwae)                                                                            |
+| Importance Weighted Autoencoder (IWAE) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/clementchadebec/benchmark_VAE/blob/main/examples/notebooks/models_training/iwae_training.ipynb) | [link](https://arxiv.org/abs/1509.00519v4)  | [link](https://github.com/yburda/iwae)  
+| Multiply Importance Weighted Autoencoder (MIWAE) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/clementchadebec/benchmark_VAE/blob/main/examples/notebooks/models_training/miwae_training.ipynb) | [link](https://arxiv.org/abs/1802.04537)  |       
+| Partially Importance Weighted Autoencoder (PIWAE) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/clementchadebec/benchmark_VAE/blob/main/examples/notebooks/models_training/piwae_training.ipynb) | [link](https://arxiv.org/abs/1802.04537)  |       
+| Combination Importance Weighted Autoencoder (CIWAE) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/clementchadebec/benchmark_VAE/blob/main/examples/notebooks/models_training/ciwae_training.ipynb) | [link](https://arxiv.org/abs/1802.04537)  |                                                                             |
 | VAE with perceptual metric similarity (MSSSIM_VAE)      | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/clementchadebec/benchmark_VAE/blob/main/examples/notebooks/models_training/ms_ssim_vae_training.ipynb) | [link](https://arxiv.org/abs/1511.06409)  |
 | Wasserstein Autoencoder (WAE)      | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/clementchadebec/benchmark_VAE/blob/main/examples/notebooks/models_training/wae_training.ipynb) | [link](https://arxiv.org/abs/1711.01558) | [link](https://github.com/tolstikhin/wae)                                  |
 | Info Variational Autoencoder (INFOVAE_MMD)      | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/clementchadebec/benchmark_VAE/blob/main/examples/notebooks/models_training/info_vae_training.ipynb) | [link](https://arxiv.org/abs/1706.02262) |                                   |

diff --git a/docs/source/models/autoencoders/ciwae.rst b/docs/source/models/autoencoders/ciwae.rst
@@ -0,0 +1,12 @@
+**********************************
+CIWAE
+**********************************
+
+.. automodule::
+   pythae.models.ciwae
+
+.. autoclass:: pythae.models.CIWAEConfig
+   :members:
+
+.. autoclass:: pythae.models.CIWAE
+   :members:
diff --git a/docs/source/models/autoencoders/miwae.rst b/docs/source/models/autoencoders/miwae.rst
@@ -0,0 +1,12 @@
+**********************************
+MIWAE
+**********************************
+
+.. automodule::
+   pythae.models.miwae
+
+.. autoclass:: pythae.models.MIWAEConfig
+   :members:
+
+.. autoclass:: pythae.models.MIWAE
+   :members:
diff --git a/docs/source/models/autoencoders/models.rst b/docs/source/models/autoencoders/models.rst
@@ -17,6 +17,9 @@ Autoencoders
     factorvae
     betatcvae
     iwae
+    ciwae
+    miwae
+    piwae
     msssimvae
     wae
     infovae
@@ -49,6 +52,9 @@ Available Models
     ~pythae.models.FactorVAE
     ~pythae.models.BetaTCVAE
     ~pythae.models.IWAE
+    ~pythae.models.CIWAE
+    ~pythae.models.MIWAE
+    ~pythae.models.PIWAE
     ~pythae.models.MSSSIM_VAE
     ~pythae.models.WAE_MMD
     ~pythae.models.INFOVAE_MMD

diff --git a/docs/source/models/autoencoders/piwae.rst b/docs/source/models/autoencoders/piwae.rst
@@ -0,0 +1,12 @@
+**********************************
+PIWAE
+**********************************
+
+.. automodule::
+   pythae.models.piwae
+
+.. autoclass:: pythae.models.PIWAEConfig
+   :members:
+
+.. autoclass:: pythae.models.PIWAE
+   :members:
diff --git a/docs/source/models/pythae.models.rst b/docs/source/models/pythae.models.rst
@@ -31,6 +31,9 @@ Available Autoencoders
     ~pythae.models.FactorVAE
     ~pythae.models.BetaTCVAE
     ~pythae.models.IWAE
+    ~pythae.models.CIWAE
+    ~pythae.models.MIWAE
+    ~pythae.models.PIWAE
     ~pythae.models.MSSSIM_VAE
     ~pythae.models.WAE_MMD
     ~pythae.models.INFOVAE_MMD

diff --git a/examples/notebooks/custom_dataset.ipynb b/examples/notebooks/custom_dataset.ipynb
@@ -16,7 +16,7 @@
    "outputs": [],
    "source": [
     "# Install the library\n",
-    "%pip install git+https://github.com/clementchadebec/benchmark_VAE.git"
+    "%pip install pythae"
    ]
   },
   {

diff --git a/examples/notebooks/mlflow_experiment_monitoring.ipynb b/examples/notebooks/mlflow_experiment_monitoring.ipynb
@@ -16,7 +16,7 @@
    "outputs": [],
    "source": [
     "# Install the library\n",
-    "%pip install git+https://github.com/clementchadebec/benchmark_VAE.git"
+    "%pip install pythae"
    ]
   },
   {