Updated BetaTCVAE results

WuLingyu · Feb 28, 2020 · 4510762 · 4510762
1 parent a36d705
commit 4510762
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 6 deletions.
diff --git a/assets/BetaTCVAE_49.png b/assets/BetaTCVAE_49.png
diff --git a/assets/recons_BetaTCVAE_49.png b/assets/recons_BetaTCVAE_49.png
diff --git a/configs/dip_vae.yaml b/configs/dip_vae.yaml
@@ -2,8 +2,8 @@ model_params:
   name: 'DIPVAE'
   in_channels: 3
   latent_dim: 128
-  lambda_diag: 10.
-  lambda_offdiag: 5.
+  lambda_diag: 4.
+  lambda_offdiag: 2.
 
 
 exp_params:

diff --git a/models/dip_vae.py b/models/dip_vae.py
@@ -137,8 +137,8 @@ def loss_function(self,
         mu = args[2]
         log_var = args[3]
 
-        kld_weight = 1. #kwargs['M_N'] # Account for the minibatch samples from the dataset
-        recons_loss =F.mse_loss(recons, input)
+        kld_weight = 1 #* kwargs['M_N'] # Account for the minibatch samples from the dataset
+        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)
@@ -150,8 +150,8 @@ def loss_function(self,
 
         cov_diag = torch.diag(cov_z) # [D]
         cov_offdiag = cov_z - torch.diag(cov_diag) # [D x D]
-        dip_loss = self.lambda_offdiag * torch.mean(cov_offdiag ** 2) + \
-                   self.lambda_diag * torch.mean((cov_diag - 1) ** 2)
+        dip_loss = self.lambda_offdiag * torch.sum(cov_offdiag ** 2) + \
+                   self.lambda_diag * torch.sum((cov_diag - 1) ** 2)
 
         loss = recons_loss + kld_weight * kld_loss + dip_loss
         return {'loss': loss,