From e0f26677d67970c4f7f804306b32ef2c0cada50f Mon Sep 17 00:00:00 2001 From: Phil Wang Date: Thu, 12 May 2022 11:12:34 -0700 Subject: [PATCH] make sure predicted mean is actually detached for all of the kl loss calculations --- denoising_diffusion_pytorch/learned_gaussian_diffusion.py | 6 ++++-- setup.py | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/denoising_diffusion_pytorch/learned_gaussian_diffusion.py b/denoising_diffusion_pytorch/learned_gaussian_diffusion.py index 45c040e..e215da8 100644 --- a/denoising_diffusion_pytorch/learned_gaussian_diffusion.py +++ b/denoising_diffusion_pytorch/learned_gaussian_diffusion.py @@ -130,10 +130,12 @@ class LearnedGaussianDiffusion(GaussianDiffusion): # kl loss with detached model predicted mean, for stability reasons as in paper - kl = normal_kl(true_mean, true_log_variance_clipped, model_mean.detach(), model_log_variance) + detached_model_mean = model_mean.detach() + + kl = normal_kl(true_mean, true_log_variance_clipped, detached_model_mean, model_log_variance) kl = meanflat(kl) * NAT - decoder_nll = -discretized_gaussian_log_likelihood(x_start, means = model_mean, log_scales = 0.5 * model_log_variance) + decoder_nll = -discretized_gaussian_log_likelihood(x_start, means = detached_model_mean, log_scales = 0.5 * model_log_variance) decoder_nll = meanflat(decoder_nll) * NAT # at the first timestep return the decoder NLL, otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t)) diff --git a/setup.py b/setup.py index 2b33658..b8bb448 100644 --- a/setup.py +++ b/setup.py @@ -3,7 +3,7 @@ from setuptools import setup, find_packages setup( name = 'denoising-diffusion-pytorch', packages = find_packages(), - version = '0.14.1', + version = '0.14.2', license='MIT', description = 'Denoising Diffusion Probabilistic Models - Pytorch', author = 'Phil Wang',