Updated results

README.md

@@ -232,6 +232,8 @@ doesn't seem to work well.
 [31]: https://github.com/AntixK/PyTorch-VAE/blob/master/assets/recons_VQVAE_29.png
 [33]: https://github.com/AntixK/PyTorch-VAE/blob/master/assets/BetaTCVAE_49.png
 [34]: https://github.com/AntixK/PyTorch-VAE/blob/master/assets/recons_BetaTCVAE_49.png
+[35]: https://github.com/AntixK/PyTorch-VAE/blob/master/assets/DIPVAE_83.png
+[36]: https://github.com/AntixK/PyTorch-VAE/blob/master/assets/recons_BDIPVAE_83.png
 
 [python-image]: https://img.shields.io/badge/Python-3.5-ff69b4.svg
 [python-url]: https://www.python.org/

configs/dip_vae.yaml

@@ -2,8 +2,8 @@ model_params:
   name: 'DIPVAE'
   in_channels: 3
   latent_dim: 128
-  lambda_diag: 4.
-  lambda_offdiag: 2.
+  lambda_diag: 0.05
+  lambda_offdiag: 0.1
 
 
 exp_params:
@@ -13,7 +13,7 @@ exp_params:
   batch_size: 144 # Better to have a square number
   LR: 0.001
   weight_decay: 0.0
-#  scheduler_gamma: 0.99
+  scheduler_gamma: 0.97
 
 trainer_params:
   gpus: 1

models/dip_vae.py

@@ -141,12 +141,16 @@ def loss_function(self,
         recons_loss =F.mse_loss(recons, input, reduction='sum')
 
 
-        kld_loss = torch.mean(-0.5 * torch.sum(1 + log_var - mu ** 2 - log_var.exp(), dim = 1), dim = 0)
+        kld_loss = torch.sum(-0.5 * torch.sum(1 + log_var - mu ** 2 - log_var.exp(), dim = 1), dim = 0)
 
         # DIP Loss
         centered_mu = mu - mu.mean(dim=1, keepdim = True) # [B x D]
         cov_mu = centered_mu.t().matmul(centered_mu).squeeze() # [D X D]
-        cov_z = cov_mu + torch.mean(torch.diagonal((2 * log_var).exp(), dim1 = 0), dim = 0) # [D x D]
+
+        # Add Variance for DIP Loss II
+        cov_z = cov_mu + torch.mean(torch.diagonal((2. * log_var).exp(), dim1 = 0), dim = 0) # [D x D]
+        # For DIp Loss I
+        # cov_z = cov_mu
 
         cov_diag = torch.diag(cov_z) # [D]
         cov_offdiag = cov_z - torch.diag(cov_diag) # [D x D]