PoincaréVAE (#51)

* [WIP] work on pvae * [WIP] fix Hyperbolic geometry and distributions * [WIP] add tests and work on PVAE * increase coverage * [WIP] fix device * [WIP] work on PVAE * [WIP] tesing the model * test * test * test * fix pvae * work on repro * change data processing in repro experiment * Add PoincaréDisk Sampler * add sampler tests * fix device issue * add tutorials and fix device setting * add PVAE to doc * minor change in sampler * minor change * small change in docs * work on reproducibility * add ref to readme * clean up readme * add Wrapped PVAE results * black and isort * update README * remove not needed * update repro * isort
clementchadebec · Sep 3, 2022 · 32b76a2 · 32b76a2
1 parent c9f0fff
commit 32b76a2
Show file tree

Hide file tree

Showing 50 changed files with 3,425 additions and 72 deletions.
diff --git a/.gitignore b/.gitignore
@@ -51,6 +51,8 @@ examples/notebooks/my_model_with_custom_archi/
 examples/notebooks/my_model/
 examples/net.py
 examples/scripts/configs/*
+examples/scripts/reproducibility/reproducibility/*
+
 
 
 

diff --git a/README.md b/README.md
@@ -36,6 +36,15 @@ provides the possibility to perform benchmark experiments and comparisons by tra
 the models with the same autoencoding neural network architecture. The feature *make your own autoencoder* 
 allows you to train any of these models with your own data and own Encoder and Decoder neural networks. It integrates an experiment monitoring tool [wandb](https://wandb.ai/) 🧪 and allows model sharing and loading from the [HuggingFace Hub](https://huggingface.co/models) 🤗 in a few lines of code.
 
+## Quick access:
+- [Installation](#installation)
+- [Implemented models](#available-models) / [Implemented samplers](#available-samplers)
+- [Reproducibility statement](#reproducibility) / [Results flavor](#results)
+- [Model training](#launching-a-model-training) / [Data generation](#launching-data-generation) / [Custom network architectures](#define-you-own-autoencoder-architecture)
+- [Model sharing with 🤗 Hub](#sharing-your-models-with-the-huggingface-hub-) / [Experiment tracking with `wandb`](#monitoring-your-experiments-with-wandb-)
+- [Tutorials](#getting-your-hands-on-the-code) / [Documentation](https://pythae.readthedocs.io/en/latest/)
+- [Contributing 🚀](#contributing-) / [Issues 🛠️](#dealing-with-issues-%EF%B8%8F)
+- [Citing this repository](#citation)
 
 # Installation
 
@@ -81,7 +90,8 @@ VAE with Inverse Autoregressive Flows (VAE_IAF) | [![Open In Colab](https://col
 | 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) | |
 | VAMP Autoencoder (VAMP) | [![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/vamp_training.ipynb) | [link](https://arxiv.org/abs/1705.07120) | [link](https://github.com/jmtomczak/vae_vampprior) |
-| Hyperspherical VAE (SVAE) | [![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/svae_training.ipynb) | [link](https://arxiv.org/abs/1804.00891) | [link](https://github.com/nicola-decao/s-vae-pytorch) |
+| Hyperspherical VAE (SVAE) | [![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/svae_training.ipynb) | [link](https://arxiv.org/abs/1804.00891) | [link](https://github.com/nicola-decao/s-vae-pytorch)
+| Poincaré Disk VAE (PoincareVAE) | [![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/pvae_training.ipynb) | [link](https://arxiv.org/abs/1901.06033) | [link](https://github.com/emilemathieu/pvae) |
 | Adversarial Autoencoder (Adversarial_AE) | [![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/adversarial_ae_training.ipynb) | [link](https://arxiv.org/abs/1511.05644)
 | Variational Autoencoder GAN (VAEGAN) 🥗 | [![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/vaegan_training.ipynb) | [link](https://arxiv.org/abs/1512.09300) | [link](https://github.com/andersbll/autoencoding_beyond_pixels)| [link](https://arxiv.org/abs/1512.09300) | [link](https://github.com/andersbll/autoencoding_beyond_pixels)
 | Vector Quantized VAE (VQVAE) | [![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/vqvae_training.ipynb) | [link](https://arxiv.org/abs/1711.00937) | [link](https://github.com/deepmind/sonnet/blob/v2/sonnet/)
@@ -102,6 +112,7 @@ Below is the list of the models currently implemented in the library.
 | Gaussian mixture (GaussianMixtureSampler) | all models | [link](https://arxiv.org/abs/1903.12436) | [link](https://github.com/ParthaEth/Regularized_autoencoders-RAE-/tree/master/models/rae) |
 | Two stage VAE sampler (TwoStageVAESampler) | all VAE based models| [link](https://openreview.net/pdf?id=B1e0X3C9tQ) | [link](https://github.com/daib13/TwoStageVAE/) |)
 | Unit sphere uniform sampler (HypersphereUniformSampler) | SVAE | [link](https://arxiv.org/abs/1804.00891) | [link](https://github.com/nicola-decao/s-vae-pytorch)
+| Poincaré Disk sampler (PoincareDiskSampler) | PoincareVAE | [link](https://arxiv.org/abs/1901.06033) | [link](https://github.com/emilemathieu/pvae)
 | 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) |

diff --git a/docs/source/models/autoencoders/models.rst b/docs/source/models/autoencoders/models.rst
@@ -22,6 +22,7 @@ Autoencoders
  infovae
  vamp
  svae
+ pvae
  aae
  vaegan
  vqvae
@@ -53,6 +54,7 @@ Available Models
  ~pythae.models.INFOVAE_MMD
  ~pythae.models.VAMP
  ~pythae.models.SVAE
+ ~pythae.models.PoincareVAE
  ~pythae.models.Adversarial_AE
  ~pythae.models.VAEGAN
  ~pythae.models.VQVAE

diff --git a/docs/source/models/autoencoders/pvae.rst b/docs/source/models/autoencoders/pvae.rst
@@ -0,0 +1,13 @@
+**********************************
+PoincareVAE
+**********************************
+
+
+.. automodule::
+ pythae.models.pvae
+
+.. autoclass:: pythae.models.PoincareVAEConfig
+ :members:
+
+.. autoclass:: pythae.models.PoincareVAE
+ :members:
diff --git a/docs/source/models/pythae.models.rst b/docs/source/models/pythae.models.rst
@@ -36,6 +36,7 @@ Available Autoencoders
  ~pythae.models.INFOVAE_MMD
  ~pythae.models.VAMP
  ~pythae.models.SVAE
+ ~pythae.models.PoincareVAE
  ~pythae.models.Adversarial_AE
  ~pythae.models.VAEGAN
  ~pythae.models.VQVAE

diff --git a/docs/source/samplers/poincare_disk_sampler.rst b/docs/source/samplers/poincare_disk_sampler.rst
@@ -0,0 +1,9 @@
+**********************************
+PoincareDiskSampler
+**********************************
+
+.. automodule::
+ pythae.samplers.pvae_sampler
+
+.. autoclass:: pythae.samplers.PoincareDiskSampler
+ :members:
diff --git a/docs/source/samplers/pythae.samplers.rst b/docs/source/samplers/pythae.samplers.rst
@@ -11,6 +11,7 @@ Samplers
  gmm_sampler
  twostage_sampler
  unit_sphere_unif_sampler
+ poincare_disk_sampler
  vamp_sampler
  rhvae_sampler
  maf_sampler
@@ -28,6 +29,7 @@ Samplers
  ~pythae.samplers.GaussianMixtureSampler
  ~pythae.samplers.TwoStageVAESampler
  ~pythae.samplers.HypersphereUniformSampler
+ ~pythae.samplers.PoincareDiskSampler
  ~pythae.samplers.VAMPSampler
  ~pythae.samplers.RHVAESampler
  ~pythae.samplers.MAFSampler

diff --git a/examples/notebooks/models_training/adversarial_ae_training.ipynb b/examples/notebooks/models_training/adversarial_ae_training.ipynb
@@ -16,8 +16,11 @@
  "metadata": {},
  "outputs": [],
  "source": [
+ "import torch\n",
  "import torchvision.datasets as datasets\n",
  "\n",
+ "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+ "\n",
  "%load_ext autoreload\n",
  "%autoreload 2"
  ]
@@ -253,7 +256,7 @@
  "metadata": {},
  "outputs": [],
  "source": [
- "reconstructions = model.reconstruct(eval_dataset[:25]).detach().cpu()"
+ "reconstructions = trained_model.reconstruct(eval_dataset[:25].to(device)).detach().cpu()"
  ]
  },
  {
@@ -301,7 +304,7 @@
  "metadata": {},
  "outputs": [],
  "source": [
- "interpolations = model.interpolate(eval_dataset[:5], eval_dataset[5:10], granularity=10).detach().cpu()"
+ "interpolations = trained_model.interpolate(eval_dataset[:5].to(device), eval_dataset[5:10].to(device), granularity=10).detach().cpu()"
  ]
  },
  {

diff --git a/examples/notebooks/models_training/ae_training.ipynb b/examples/notebooks/models_training/ae_training.ipynb
@@ -16,8 +16,11 @@
  "metadata": {},
  "outputs": [],
  "source": [
+ "import torch\n",
  "import torchvision.datasets as datasets\n",
  "\n",
+ "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+ "\n",
  "%load_ext autoreload\n",
  "%autoreload 2"
  ]
@@ -252,7 +255,7 @@
  "metadata": {},
  "outputs": [],
  "source": [
- "reconstructions = model.reconstruct(eval_dataset[:25]).detach().cpu()"
+ "reconstructions = trained_model.reconstruct(eval_dataset[:25].to(device)).detach().cpu()"
  ]
  },
  {
@@ -300,7 +303,7 @@
  "metadata": {},
  "outputs": [],
  "source": [
- "interpolations = model.interpolate(eval_dataset[:5], eval_dataset[5:10], granularity=10).detach().cpu()"
+ "interpolations = trained_model.interpolate(eval_dataset[:5].to(device), eval_dataset[5:10].to(device), granularity=10).detach().cpu()"
  ]
  },
  {

diff --git a/examples/notebooks/models_training/beta_tc_vae_training.ipynb b/examples/notebooks/models_training/beta_tc_vae_training.ipynb
@@ -16,8 +16,11 @@
  "metadata": {},
  "outputs": [],
  "source": [
+ "import torch\n",
  "import torchvision.datasets as datasets\n",
  "\n",
+ "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+ "\n",
  "%load_ext autoreload\n",
  "%autoreload 2"
  ]
@@ -256,7 +259,7 @@
  "metadata": {},
  "outputs": [],
  "source": [
- "reconstructions = model.reconstruct(eval_dataset[:25]).detach().cpu()"
+ "reconstructions = trained_model.reconstruct(eval_dataset[:25].to(device)).detach().cpu()"
  ]
  },
  {
@@ -304,7 +307,7 @@
  "metadata": {},
  "outputs": [],
  "source": [
- "interpolations = model.interpolate(eval_dataset[:5], eval_dataset[5:10], granularity=10).detach().cpu()"
+ "interpolations = trained_model.interpolate(eval_dataset[:5].to(device), eval_dataset[5:10].to(device), granularity=10).detach().cpu()"
  ]
  },
  {

diff --git a/examples/notebooks/models_training/beta_vae_training.ipynb b/examples/notebooks/models_training/beta_vae_training.ipynb
@@ -16,8 +16,11 @@
  "metadata": {},
  "outputs": [],
  "source": [
+ "import torch\n",
  "import torchvision.datasets as datasets\n",
  "\n",
+ "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+ "\n",
  "%load_ext autoreload\n",
  "%autoreload 2"
  ]
@@ -254,7 +257,7 @@
  "metadata": {},
  "outputs": [],
  "source": [
- "reconstructions = model.reconstruct(eval_dataset[:25]).detach().cpu()"
+ "reconstructions = trained_model.reconstruct(eval_dataset[:25].to(device)).detach().cpu()"
  ]
  },
  {
@@ -302,7 +305,7 @@
  "metadata": {},
  "outputs": [],
  "source": [
- "interpolations = model.interpolate(eval_dataset[:5], eval_dataset[5:10], granularity=10).detach().cpu()"
+ "interpolations = trained_model.interpolate(eval_dataset[:5].to(device), eval_dataset[5:10].to(device), granularity=10).detach().cpu()"
  ]
  },
  {

diff --git a/examples/notebooks/models_training/disentangled_beta_vae_training.ipynb b/examples/notebooks/models_training/disentangled_beta_vae_training.ipynb
@@ -16,8 +16,11 @@
  "metadata": {},
  "outputs": [],
  "source": [
+ "import torch\n",
  "import torchvision.datasets as datasets\n",
  "\n",
+ "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+ "\n",
  "%load_ext autoreload\n",
  "%autoreload 2"
  ]
@@ -256,7 +259,7 @@
  "metadata": {},
  "outputs": [],
  "source": [
- "reconstructions = model.reconstruct(eval_dataset[:25]).detach().cpu()"
+ "reconstructions = trained_model.reconstruct(eval_dataset[:25].to(device)).detach().cpu()"
  ]
  },
  {
@@ -304,7 +307,7 @@
  "metadata": {},
  "outputs": [],
  "source": [
- "interpolations = model.interpolate(eval_dataset[:5], eval_dataset[5:10], granularity=10).detach().cpu()"
+ "interpolations = trained_model.interpolate(eval_dataset[:5].to(device), eval_dataset[5:10].to(device), granularity=10).detach().cpu()"
  ]
  },
  {

diff --git a/examples/notebooks/models_training/factor_vae_training.ipynb b/examples/notebooks/models_training/factor_vae_training.ipynb
@@ -16,8 +16,11 @@
  "metadata": {},
  "outputs": [],
  "source": [
+ "import torch\n",
  "import torchvision.datasets as datasets\n",
  "\n",
+ "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+ "\n",
  "%load_ext autoreload\n",
  "%autoreload 2"
  ]
@@ -254,7 +257,7 @@
  "metadata": {},
  "outputs": [],
  "source": [
- "reconstructions = model.reconstruct(eval_dataset[:25]).detach().cpu()"
+ "reconstructions = trained_model.reconstruct(eval_dataset[:25].to(device)).detach().cpu()"
  ]
  },
  {
@@ -302,7 +305,7 @@
  "metadata": {},
  "outputs": [],
  "source": [
- "interpolations = model.interpolate(eval_dataset[:5], eval_dataset[5:10], granularity=10).detach().cpu()"
+ "interpolations = trained_model.interpolate(eval_dataset[:5].to(device), eval_dataset[5:10].to(device), granularity=10).detach().cpu()"
  ]
  },
  {
@@ -340,7 +343,7 @@
  "name": "python",
  "nbconvert_exporter": "python",
  "pygments_lexer": "ipython3",
- "version": "3.8.13"
+ "version": "3.8.12"
  },
  "orig_nbformat": 4
  },

diff --git a/examples/notebooks/models_training/hvae_training.ipynb b/examples/notebooks/models_training/hvae_training.ipynb
@@ -16,8 +16,11 @@
  "metadata": {},
  "outputs": [],
  "source": [
+ "import torch\n",
  "import torchvision.datasets as datasets\n",
  "\n",
+ "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+ "\n",
  "%load_ext autoreload\n",
  "%autoreload 2"
  ]
@@ -255,7 +258,7 @@
  "metadata": {},
  "outputs": [],
  "source": [
- "reconstructions = model.reconstruct(eval_dataset[:25]).detach().cpu()"
+ "reconstructions = trained_model.reconstruct(eval_dataset[:25].to(device)).detach().cpu()"
  ]
  },
  {
@@ -303,7 +306,7 @@
  "metadata": {},
  "outputs": [],
  "source": [
- "interpolations = model.interpolate(eval_dataset[:5], eval_dataset[5:10], granularity=10).detach().cpu()"
+ "interpolations = trained_model.interpolate(eval_dataset[:5].to(device), eval_dataset[5:10].to(device), granularity=10).detach().cpu()"
  ]
  },
  {
@@ -341,7 +344,7 @@
  "name": "python",
  "nbconvert_exporter": "python",
  "pygments_lexer": "ipython3",
- "version": "3.8.11"
+ "version": "3.8.12"
  },
  "orig_nbformat": 4
  },

diff --git a/examples/notebooks/models_training/info_vae_training.ipynb b/examples/notebooks/models_training/info_vae_training.ipynb
@@ -16,8 +16,11 @@
  "metadata": {},
  "outputs": [],
  "source": [
+ "import torch\n",
  "import torchvision.datasets as datasets\n",
  "\n",
+ "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+ "\n",
  "%load_ext autoreload\n",
  "%autoreload 2"
  ]
@@ -263,7 +266,7 @@
  "metadata": {},
  "outputs": [],
  "source": [
- "reconstructions = model.reconstruct(eval_dataset[:25]).detach().cpu()"
+ "reconstructions = trained_model.reconstruct(eval_dataset[:25].to(device)).detach().cpu()"
  ]
  },
  {
@@ -311,7 +314,7 @@
  "metadata": {},
  "outputs": [],
  "source": [
- "interpolations = model.interpolate(eval_dataset[:5], eval_dataset[5:10], granularity=10).detach().cpu()"
+ "interpolations = trained_model.interpolate(eval_dataset[:5].to(device), eval_dataset[5:10].to(device), granularity=10).detach().cpu()"
  ]
  },
  {
@@ -349,7 +352,7 @@
  "name": "python",
  "nbconvert_exporter": "python",
  "pygments_lexer": "ipython3",
- "version": "3.8.11"
+ "version": "3.8.12"
  },
  "orig_nbformat": 4
  },