Added Joint VAE

JasonTian0323 · Jan 27, 2020 · d282f67 · d282f67
1 parent 7f9df52
commit d282f67
Show file tree

Hide file tree

Showing 10 changed files with 364 additions and 14 deletions.
diff --git a/README.md b/README.md
@@ -88,7 +88,8 @@ logging_params:
 | IWAE (5 Samples)      |[Link](https://arxiv.org/abs/1804.03599)          |    ![][10]    | ![][9]  |
 | DFCVAE                |[Link](https://arxiv.org/abs/1610.00291)          |    ![][12]    | ![][11] |
 | MSSIM VAE             |[Link](https://arxiv.org/abs/1511.06409)          |    ![][14]    | ![][13] |
-| Categorical VAE (CIFAR10)|[Link](https://arxiv.org/abs/1611.01144)       |    ![][18]    | ![][17] |
+| Categorical VAE       |[Link](https://arxiv.org/abs/1611.01144)          |    ![][18]    | ![][17] |
+| Joint VAE             |[Link](https://arxiv.org/abs/1804.00104)          |    ![][20]    | ![][19] |
 
 <!-- | Gamma VAE             |[Link](https://arxiv.org/abs/1610.05683)          |    ![][16]    | ![][15] |-->
 <!--| Disentangled Beta-VAE |[Link](https://arxiv.org/abs/1804.03599)          |    ![][10]     | ![][9] |-->
@@ -104,6 +105,7 @@ logging_params:
 - [x] WAE-MMD
 - [x] Conditional VAE
 - [x] Categorical VAE (Gumbel-Softmax VAE)
+- [x] Joint VAE
 - [ ] Gamma VAE (in progress)
 - [ ] Beta TC-VAE (in progress) 
 - [ ] Vamp VAE (in progress)
@@ -117,7 +119,7 @@ logging_params:
 - [ ] VQVAE
 - [ ] StyleVAE
 - [ ] Sequential VAE
-- [ ] Joint VAE
+
 
 ### Contributing
 If you have trained a better model, using these implementations, by fine-tuning the hyper-params in the config file,
@@ -154,7 +156,8 @@ I would be happy to include your result (along with your config file) in this re
 [16]: https://github.com/AntixK/PyTorch-VAE/blob/master/assets/recons_ConditionalVAE_20.png
 [17]: https://github.com/AntixK/PyTorch-VAE/blob/master/assets/CategoricalVAE_20.png
 [18]: https://github.com/AntixK/PyTorch-VAE/blob/master/assets/recons_CategoricalVAE_20.png
-
+[19]: https://github.com/AntixK/PyTorch-VAE/blob/master/assets/JointVAE_20.png
+[20]: https://github.com/AntixK/PyTorch-VAE/blob/master/assets/recons_JointVAE_20.png
 
 [python-image]: https://img.shields.io/badge/Python-3.5-ff69b4.svg
 [python-url]: https://www.python.org/

diff --git a/configs/cat_vae.yaml b/configs/cat_vae.yaml
@@ -1,14 +1,15 @@
 model_params:
   name: 'CategoricalVAE'
   in_channels: 3
-  latent_dim: 256
-  categorical_dim: 10 # Equal to Num classes
+  latent_dim: 512
+  categorical_dim: 40
   temperature: 0.5
   anneal_rate: 0.00003
   anneal_interval: 100
+  alpha: 8.0
 
 exp_params:
-  dataset: cifar10
+  dataset: celeba
   data_path: "../../shared/Data/"
   img_size: 64
   batch_size: 144 # Better to have a square number
@@ -17,9 +18,9 @@ exp_params:
   scheduler_gamma: 0.95
 
 trainer_params:
-  gpus: 1
+  gpus: [1]
   max_nb_epochs: 50
-  max_epochs: 250
+  max_epochs: 50
 
 logging_params:
   save_dir: "logs/"

diff --git a/configs/joint_vae.yaml b/configs/joint_vae.yaml
@@ -0,0 +1,36 @@
+model_params:
+  name: 'JointVAE'
+  in_channels: 3
+  latent_dim: 512
+  categorical_dim: 40
+  latent_min_capacity: 0.0
+  latent_max_capacity: 20.0
+  latent_gamma: 30.
+  latent_num_iter: 25000
+  categorical_min_capacity: 0.0
+  categorical_max_capacity: 20.0
+  categorical_gamma: 30.
+  categorical_num_iter: 25000
+  temperature: 0.5
+  anneal_rate: 0.00003
+  anneal_interval: 100
+  alpha: 20.0
+
+exp_params:
+  dataset: celeba
+  data_path: "../../shared/Data/"
+  img_size: 64
+  batch_size: 144 # Better to have a square number
+  LR: 0.005
+  weight_decay: 0.0
+  scheduler_gamma: 0.95
+
+trainer_params:
+  gpus: [1]
+  max_nb_epochs: 50
+  max_epochs: 50
+
+logging_params:
+  save_dir: "logs/"
+  name: "JointVAE"
+  manual_seed: 1265
diff --git a/experiment.py b/experiment.py
@@ -205,7 +205,7 @@ def data_transforms(self):
                                             transforms.ToTensor(),
                                             SetRange])
 
-        if self.params['dataset'] == 'cifar10':
+        elif self.params['dataset'] == 'cifar10':
             transform = transforms.Compose([transforms.RandomHorizontalFlip(),
                                             transforms.ToTensor(),
                                             transforms.Lambda(lambda img:

diff --git a/models/__init__.py b/models/__init__.py
@@ -11,12 +11,13 @@
 from .mssim_vae import MSSIMVAE
 from .fvae import *
 from .cat_vae import *
+from .joint_vae import *
 
 # Aliases
 VAE = VanillaVAE
 GaussianVAE = VanillaVAE
 CVAE = ConditionalVAE
-GUMBELVAE = CategoricalVAE
+GumbelVAE = CategoricalVAE
 
 vae_models = {'VanillaVAE':VanillaVAE,
               'WAE_MMD':WAE_MMD,
@@ -29,4 +30,5 @@
               'DFCVAE':DFCVAE,
               'MSSIMVAE':MSSIMVAE,
               'FactorVAE':FactorVAE,
-              'CategoricalVAE':CategoricalVAE}
+              'CategoricalVAE':CategoricalVAE,
+              'JointVAE':JointVAE}
diff --git a/models/__pycache__/__init__.cpython-36.pyc b/models/__pycache__/__init__.cpython-36.pyc
diff --git a/models/cat_vae.py b/models/cat_vae.py
@@ -16,6 +16,7 @@ def __init__(self,
                  temperature: float = 0.5,
                  anneal_rate: float = 3e-5,
                  anneal_interval: int = 100, # every 100 batches
+                 alpha: float = 30.,
                  **kwargs) -> None:
         super(CategoricalVAE, self).__init__()
 
@@ -25,6 +26,7 @@ def __init__(self,
         self.min_temp = temperature
         self.anneal_rate = anneal_rate
         self.anneal_interval = anneal_interval
+        self.alpha = alpha
 
         modules = []
         if hidden_dims is None:
@@ -171,7 +173,7 @@ def loss_function(self,
         kld_loss = torch.mean(torch.sum(h1 - h2, dim =(1,2)), dim=0)
 
         # kld_weight = 1.2
-        loss = 30. * recons_loss + kld_weight * kld_loss
+        loss = self.alpha * recons_loss + kld_weight * kld_loss
         return {'loss': loss, 'Reconstruction_Loss':recons_loss, 'KLD':-kld_loss}
 
     def sample(self,