Added support for latest pytorch lightning version

Author: AntixK

.gitignore

@@ -1,5 +1,5 @@
 
-data/
+Data/
 logs/
 
 VanillaVAE/version_0/

README.md

@@ -39,6 +39,7 @@ $ cd PyTorch-VAE
 $ python run.py -c configs/<config-file-name.yaml>
 ```
 **Config file template**
+
 ```yaml
 model_params:
   name: "<name of VAE model>"
@@ -48,10 +49,15 @@ model_params:
     .
     .
 
-exp_params:
+data_params:
   data_path: "<path to the celebA dataset>"
-  img_size: 64    # Models are designed to work for this size
-  batch_size: 64  # Better to have a square number
+  train_batch_size: 64 # Better to have a square number
+  val_batch_size:  64
+  patch_size: 64  # Models are designed to work for this size
+  num_workers: 4
+  
+exp_params:
+  manual_seed: 1265
   LR: 0.005
   weight_decay:
     .         # Other arguments required for training, like scheduler etc.
@@ -60,7 +66,7 @@ exp_params:
 
 trainer_params:
   gpus: 1         
-  max_nb_epochs: 50
+  max_epochs: 100
   gradient_clip_val: 1.5
     .
     .
@@ -69,15 +75,17 @@ trainer_params:
 logging_params:
   save_dir: "logs/"
   name: "<experiment name>"
-  manual_seed: 
 ```
 
 **View TensorBoard Logs**
 ```
 $ cd logs/<experiment name>/version_<the version you want>
-$ tensorboard --logdir tf
+$ tensorboard --logdir .
 ```
 
+**Note:** The default dataset is CelebA. However, there has been many issues with downloading the dataset from google drive (owing to some file structure changes). So, the recommendation is to download the [file](https://drive.google.com/file/d/1m8-EBPgi5MRubrm6iQjafK2QMHDBMSfJ/view?usp=sharing) from google drive directly and extract to the path of your choice. The default path assumed in the config files is `Data/celeba/img_align_celeba'. But you can change it acording to your preference.
+
+
 ----
 <h2 align="center">
   <b>Results</b><br>

configs/bbvae.yaml

@@ -7,21 +7,25 @@ model_params:
   max_capacity: 25
   Capacity_max_iter: 10000
 
+data_params:
+  data_path: "Data/"
+  train_batch_size: 64
+  val_batch_size:  64
+  patch_size: 64
+  num_workers: 4
+  
 exp_params:
-  dataset: celeba
-  data_path: "../../shared/Data/"
-  img_size: 64
-  batch_size: 144 # Better to have a square number
-  LR: 0.0005
+  LR: 0.005
   weight_decay: 0.0
   scheduler_gamma: 0.95
+  kld_weight: 0.00025
+  manual_seed: 1265
 
 trainer_params:
-  gpus: 1
-  max_nb_epochs: 50
-  max_epochs: 50
+  gpus: [1]
+  max_epochs: 10
 
 logging_params:
   save_dir: "logs/"
-  name: "BetaVAE_B"
   manual_seed: 1265
+  name: 'BetaVAE'

configs/betatc_vae.yaml

@@ -7,21 +7,25 @@ model_params:
   beta:  6.
   gamma: 1.
 
+data_params:
+  data_path: "Data/"
+  train_batch_size: 64
+  val_batch_size:  64
+  patch_size: 64
+  num_workers: 4
+
+
 exp_params:
-  dataset: celeba
-  data_path: "../../shared/momo/Data/"
-  img_size: 64
-  batch_size: 144 # Better to have a square number
-  LR: 0.001
+  LR: 0.005
   weight_decay: 0.0
-#  scheduler_gamma: 0.99
+  scheduler_gamma: 0.95
+  kld_weight: 0.00025
+  manual_seed: 1265
 
 trainer_params:
-  gpus: 1
-  max_nb_epochs: 50
-  max_epochs: 50
+  gpus: [1]
+  max_epochs: 10
 
 logging_params:
   save_dir: "logs/"
-  name: "BetaTCVAE"
-  manual_seed: 1265
+  name: 'BetaTCVAE'

configs/bhvae.yaml

@@ -5,21 +5,25 @@ model_params:
   loss_type: 'H'
   beta: 10.
 
+data_params:
+  data_path: "Data/"
+  train_batch_size: 64
+  val_batch_size:  64
+  patch_size: 64
+  num_workers: 4
+
+
 exp_params:
-  dataset: celeba
-  data_path: "../../shared/Data/"
-  img_size: 64
-  batch_size: 144 # Better to have a square number
-  LR: 0.0005
+  LR: 0.005
   weight_decay: 0.0
   scheduler_gamma: 0.95
+  kld_weight: 0.00025
+  manual_seed: 1265
 
 trainer_params:
-  gpus: 1
-  max_nb_epochs: 50
-  max_epochs: 50
+  gpus: [1]
+  max_epochs: 10
 
 logging_params:
   save_dir: "logs/"
-  name: "BetaVAE_H"
-  manual_seed: 1265
+  name: 'BetaVAE'

configs/cat_vae.yaml

@@ -8,21 +8,25 @@ model_params:
   anneal_interval: 100
   alpha: 1.0
 
+data_params:
+  data_path: "Data/"
+  train_batch_size: 64
+  val_batch_size:  64
+  patch_size: 64
+  num_workers: 4
+
+
 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
+  kld_weight: 0.00025
+  manual_seed: 1265
 
 trainer_params:
   gpus: [1]
-  max_nb_epochs: 50
-  max_epochs: 50
+  max_epochs: 10
 
 logging_params:
   save_dir: "logs/"
   name: "CategoricalVAE"
-  manual_seed: 1265

configs/cvae.yaml

@@ -4,21 +4,25 @@ model_params:
   num_classes: 40
   latent_dim: 128
 
+data_params:
+  data_path: "Data/"
+  train_batch_size: 64
+  val_batch_size:  64
+  patch_size: 64
+  num_workers: 4
+
+
 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
+  kld_weight: 0.00025
+  manual_seed: 1265
 
 trainer_params:
-  gpus: 1
-  max_nb_epochs: 50
-  max_epochs: 50
+  gpus: [1]
+  max_epochs: 10
 
 logging_params:
   save_dir: "logs/"
-  name: "ConditionalVAE"
-  manual_seed: 1265
+  name: "ConditionalVAE"

configs/dfc_vae.yaml

@@ -3,21 +3,25 @@ model_params:
   in_channels: 3
   latent_dim: 128
 
+data_params:
+  data_path: "Data/"
+  train_batch_size: 64
+  val_batch_size:  64
+  patch_size: 64
+  num_workers: 4
+
+
 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
+  kld_weight: 0.00025
+  manual_seed: 1265
 
 trainer_params:
-  gpus: 1
-  max_nb_epochs: 50
-  max_epochs: 50
+  gpus: [1]
+  max_epochs: 10
 
 logging_params:
   save_dir: "logs/"
   name: "DFCVAE"
-  manual_seed: 1265

configs/dip_vae.yaml

@@ -6,19 +6,24 @@ model_params:
   lambda_offdiag: 0.1
 
 
+data_params:
+  data_path: "Data/"
+  train_batch_size: 64
+  val_batch_size:  64
+  patch_size: 64
+  num_workers: 4
+
+
 exp_params:
-  dataset: celeba
-  data_path: "../../shared/momo/Data/"
-  img_size: 64
-  batch_size: 144 # Better to have a square number
   LR: 0.001
   weight_decay: 0.0
   scheduler_gamma: 0.97
+  kld_weight: 1
+  manual_seed: 1265
 
 trainer_params:
-  gpus: 1
-  max_nb_epochs: 50
-  max_epochs: 50
+  gpus: [1]
+  max_epochs: 10
 
 logging_params:
   save_dir: "logs/"

configs/factorvae.yaml

@@ -4,25 +4,31 @@ model_params:
   latent_dim: 128
   gamma: 6.4
 
+data_params:
+  data_path: "Data/"
+  train_batch_size: 64
+  val_batch_size:  64
+  patch_size: 64
+  num_workers: 4
+
+
 exp_params:
-  dataset: celeba
-  data_path: "../../shared/Data/"
   submodel: 'discriminator'
   retain_first_backpass: True
-  img_size: 64
-  batch_size: 144 # Better to have a square number
   LR: 0.005
   weight_decay: 0.0
-  scheduler_gamma: 0.95
   LR_2: 0.005
   scheduler_gamma_2: 0.95
+  scheduler_gamma: 0.95
+  kld_weight: 0.00025
+  manual_seed: 1265
 
 trainer_params:
-  gpus: [3]
-  max_nb_epochs: 30
-  max_epochs: 50
+  gpus: [1]
+  max_epochs: 10
 
 logging_params:
   save_dir: "logs/"
-  name: "FactorVAE"
-  manual_seed: 1265
+  name: "FactorVAE"  
+  
+