Updated GammaVAE

README.md

@@ -52,13 +52,18 @@ exp_params:
   img_size: 64    # Models are designed to work for this size
   batch_size: 64  # Better to have a square number
   LR: 0.005
-    .         # Other arguments required for training like scheduler etc.
+  weight_decay:
+    .         # Other arguments required for training, like scheduler etc.
     .
     .
 
 trainer_params:
   gpus: 1         
   max_nb_epochs: 50
+  gradient_clip_val: 0.005
+    .
+    .
+    .
 
 logging_params:
   save_dir: "logs/"

configs/bhvae.yaml

@@ -12,6 +12,7 @@ exp_params:
   img_size: 64
   batch_size: 144 # Better to have a square number
   LR: 0.0005
+  weight_decay: 0.0
   scheduler_gamma: 0.95
 
 trainer_params:

configs/bvae.yaml

@@ -10,6 +10,7 @@ exp_params:
   img_size: 64
   batch_size: 144 # Better to have a square number
   LR: 0.0005
+  weight_decay: 0.0
   scheduler_gamma: 0.95
 
 trainer_params:

configs/cvae.yaml

@@ -0,0 +1,23 @@
+model_params:
+  name: 'ConditionalVAE'
+  in_channels: 3
+  num_classes: 40
+  latent_dim: 128
+
+exp_params:
+  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
+
+
+logging_params:
+  save_dir: "logs/"
+  name: "ConditionalVAE"
+  manual_seed: 1265

configs/dfc_vae.yaml

@@ -8,6 +8,7 @@ exp_params:
   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:

configs/factorvae.yaml

@@ -11,6 +11,7 @@ exp_params:
   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

configs/gammavae.yaml

@@ -10,12 +10,14 @@ exp_params:
   data_path: "../../shared/Data/"
   img_size: 64
   batch_size: 144 # Better to have a square number
-  LR: 0.005
-  scheduler_gamma: 0.95
+  LR: 0.003
+  weight_decay: 0.0005
+
 
 trainer_params:
   gpus: 1
   max_nb_epochs: 50
+  gradient_clip_val: 0.8
 
 
 logging_params:

configs/hvae.yaml

@@ -10,6 +10,7 @@ exp_params:
   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:

configs/iwae.yaml

@@ -8,6 +8,7 @@ exp_params:
   img_size: 64
   batch_size: 144 # Better to have a square number
   LR: 0.007
+  weight_decay: 0.0
   scheduler_gamma: 0.95
 
 trainer_params:

configs/mssim_vae.yaml

@@ -8,6 +8,7 @@ exp_params:
   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:

configs/vae.yaml

@@ -8,10 +8,11 @@ exp_params:
   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: 0
+  gpus: 1
   max_nb_epochs: 50
 
 

configs/vampvae.yaml

@@ -8,6 +8,7 @@ exp_params:
   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:

configs/wae_mmd_imq.yaml

@@ -10,6 +10,7 @@ exp_params:
   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:

configs/wae_mmd_rbf.yaml

@@ -10,6 +10,7 @@ exp_params:
   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:

experiment.py

@@ -68,22 +68,27 @@ def validation_end(self, outputs):
         return {'val_loss': avg_loss, 'log': tensorboard_logs}
 
     def sample_images(self):
-        samples = self.model.sample(self.params['batch_size'], self.curr_device).cpu()
-        vutils.save_image(samples.data,
-                          f"{self.logger.save_dir}{self.logger.name}/version_{self.logger.version}/"
-                          f"{self.logger.name}_{self.current_epoch}.png",
-                          normalize=True,
-                          nrow=int(math.sqrt(self.params['batch_size'])))
-
         # Get sample reconstruction image
-        test_input, _ = next(iter(self.sample_dataloader))
+        test_input, test_label = next(iter(self.sample_dataloader))
         test_input = test_input.to(self.curr_device)
-        recons = self.model.generate(test_input)
+        test_label = test_label.to(self.curr_device)
+        recons = self.model.generate(test_input, labels = test_label)
         vutils.save_image(recons.data,
                           f"{self.logger.save_dir}{self.logger.name}/version_{self.logger.version}/"
                           f"recons_{self.logger.name}_{self.current_epoch}.png",
                           normalize=True,
                           nrow=int(math.sqrt(self.params['batch_size'])))
+
+        samples = self.model.sample(self.params['batch_size'],
+                                    self.curr_device,
+                                    labels = test_label).cpu()
+        vutils.save_image(samples.data,
+                          f"{self.logger.save_dir}{self.logger.name}/version_{self.logger.version}/"
+                          f"{self.logger.name}_{self.current_epoch}.png",
+                          normalize=True,
+                          nrow=int(math.sqrt(self.params['batch_size'])))
+
+
         del test_input, recons, samples
 
     # def backward(self, use_amp, loss, optimizer):
@@ -97,7 +102,9 @@ def configure_optimizers(self):
         optims = []
         scheds = []
 
-        optimizer = optim.Adam(self.model.parameters(), lr=self.params['LR'])
+        optimizer = optim.Adam(self.model.parameters(),
+                               lr=self.params['LR'],
+                               weight_decay=self.params['weight_decay'])
         optims.append(optimizer)
         # Check if more than 1 optimizer is required (Used for adversarial training)
         try:

models/beta_vae.py

@@ -149,7 +149,7 @@ def loss_function(self,
         else:
             raise ValueError('Undefined loss type.')
 
-        return {'loss': loss, 'Reconstruction Loss':recons_loss, 'KLD':kld_loss}
+        return {'loss': loss, 'Reconstruction_Loss':recons_loss, 'KLD':kld_loss}
 
     def sample(self,
                num_samples:int,

models/cvae.py

@@ -117,7 +117,7 @@ def reparameterize(self, mu: Tensor, logvar: Tensor) -> Tensor:
         return eps * std + mu
 
     def forward(self, input: Tensor, **kwargs) -> List[Tensor]:
-        y = kwargs['labels']
+        y = kwargs['labels'].float()
         embedded_class = self.embed_class(y)
         embedded_class = embedded_class.view(-1, self.img_size, self.img_size).unsqueeze(1)
         embedded_input = self.embed_data(input)
@@ -144,31 +144,34 @@ def loss_function(self,
         kld_loss = torch.mean(-0.5 * torch.sum(1 + log_var - mu ** 2 - log_var.exp(), dim = 1), dim = 0)
 
         loss = recons_loss + kld_weight * kld_loss
-        return {'loss': loss, 'Reconstruction Loss':recons_loss, 'KLD':-kld_loss}
+        return {'loss': loss, 'Reconstruction_Loss':recons_loss, 'KLD':-kld_loss}
 
     def sample(self,
                num_samples:int,
-               current_device: int) -> Tensor:
+               current_device: int,
+               **kwargs) -> Tensor:
         """
         Samples from the latent space and return the corresponding
         image space map.
         :param num_samples: (Int) Number of samples
         :param current_device: (Int) Device to run the model
         :return: (Tensor)
         """
+        y = kwargs['labels'].float()
         z = torch.randn(num_samples,
                         self.latent_dim)
 
-        z = z.cuda(current_device)
+        z = z.to(current_device)
 
+        z = torch.cat([z, y], dim=1)
         samples = self.decode(z)
         return samples
 
-    def generate(self, x: Tensor) -> Tensor:
+    def generate(self, x: Tensor, **kwargs) -> Tensor:
         """
         Given an input image x, returns the reconstructed image
         :param x: (Tensor) [B x C x H x W]
         :return: (Tensor) [B x C x H x W]
         """
 
-        return self.forward(x)[0]
+        return self.forward(x, **kwargs)[0]

models/dfcvae.py

@@ -187,7 +187,7 @@ def loss_function(self,
         kld_loss = torch.mean(-0.5 * torch.sum(1 + log_var - mu ** 2 - log_var.exp(), dim = 1), dim = 0)
 
         loss = self.beta * (recons_loss + feature_loss) + self.alpha * kld_weight * kld_loss
-        return {'loss': loss, 'Reconstruction Loss':recons_loss, 'KLD':-kld_loss}
+        return {'loss': loss, 'Reconstruction_Loss':recons_loss, 'KLD':-kld_loss}
 
     def sample(self,
                num_samples:int,