From d720a8854349fceed6a564590efe34e570b60f05 Mon Sep 17 00:00:00 2001
From: Pyeong Eun Kim <35253323+PyeongKim@users.noreply.github.com>
Date: Mon, 24 May 2021 11:39:48 +0900
Subject: [PATCH] Update loss.py

Dimension of mu and logvar is (Batch x d_latent).
If we just calculate mean of all as previously implemented, this would result in mean of all d_latent regardless of individual data.
Instead, we have to make sum of latent representation, and then get mean along batch.
I think it is more close to original meaning of KL Divergence.
---
 transvae/loss.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/transvae/loss.py b/transvae/loss.py
index 37c8fae..328dc65 100644
--- a/transvae/loss.py
+++ b/transvae/loss.py
@@ -11,7 +11,7 @@ def vae_loss(x, x_out, mu, logvar, weights, beta=1):
     x = x.contiguous().view(-1)
     x_out = x_out.contiguous().view(-1, x_out.size(2))
     BCE = F.cross_entropy(x_out, x, reduction='mean', weight=weights)
-    KLD = beta * -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
+    KLD = beta * torch.mean(-0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1), dim=0)
     if torch.isnan(KLD):
         KLD = 0.
     return BCE + KLD, BCE, KLD
@@ -24,7 +24,7 @@ def trans_vae_loss(x, x_out, mu, logvar, true_len, pred_len, weights, beta=1):
     true_len = true_len.contiguous().view(-1)
     BCEmol = F.cross_entropy(x_out, x, reduction='mean', weight=weights)
     BCEmask = F.cross_entropy(pred_len, true_len, reduction='mean')
-    KLD = beta * -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
+    KLD = beta * torch.mean(-0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1), dim=0)
     if torch.isnan(KLD):
         KLD = 0.
     return BCEmol + BCEmask + KLD, BCEmol, BCEmask, KLD