Replaced KL divergence estimation with deterministic KL (#760)

Co-authored-by: Jirka <[email protected]>
Lightning-Universe · Nov 8, 2021 · 51942bc · 51942bc
1 parent 710477d
commit 51942bc
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- VAE now uses deterministic KL divergence during training, previously estimated KL divergence by random sampling ([#760](https://github.com/PyTorchLightning/lightning-bolts/pull/760))
 
 ### Deprecated
 

diff --git a/pl_bolts/models/autoencoders/basic_vae/basic_vae_module.py b/pl_bolts/models/autoencoders/basic_vae/basic_vae_module.py
@@ -131,10 +131,7 @@ def step(self, batch, batch_idx):
 
         recon_loss = F.mse_loss(x_hat, x, reduction="mean")
 
-        log_qz = q.log_prob(z)
-        log_pz = p.log_prob(z)
-
-        kl = log_qz - log_pz
+        kl = torch.distributions.kl_divergence(q, p)
         kl = kl.mean()
         kl *= self.kl_coeff
Original file line number	Diff line number	Diff line change
Expand Up		@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

		### Changed

		- VAE now uses deterministic KL divergence during training, previously estimated KL divergence by random sampling ([#760](https://github.com/PyTorchLightning/lightning-bolts/pull/760))

		### Deprecated

Expand Down