From 402b7c26df7e0c4cf009be64de781b08ead2e7ef Mon Sep 17 00:00:00 2001 From: Phil Wang Date: Tue, 10 May 2022 15:23:34 -0700 Subject: [PATCH] calculate noise schedule with float64 for numerical accuracy --- .../denoising_diffusion_pytorch.py | 32 +++++++++++-------- setup.py | 2 +- 2 files changed, 19 insertions(+), 15 deletions(-) diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py index 39da378..78d97bd 100644 --- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py +++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py @@ -316,11 +316,11 @@ def cosine_beta_schedule(timesteps, s = 0.008): as proposed in https://openreview.net/forum?id=-NEXDKk8gZ """ steps = timesteps + 1 - x = torch.linspace(0, timesteps, steps) + x = torch.linspace(0, timesteps, steps, dtype = torch.float64) alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * torch.pi * 0.5) ** 2 alphas_cumprod = alphas_cumprod / alphas_cumprod[0] betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1]) - return torch.clip(betas, 0, 0.999) + return torch.clip(betas, 0, 0.9999) class GaussianDiffusion(nn.Module): def __init__( @@ -347,17 +347,21 @@ class GaussianDiffusion(nn.Module): self.num_timesteps = int(timesteps) self.loss_type = loss_type - self.register_buffer('betas', betas) - self.register_buffer('alphas_cumprod', alphas_cumprod) - self.register_buffer('alphas_cumprod_prev', alphas_cumprod_prev) + # helper function to register buffer from float64 to float32 + + register_buffer = lambda name, val: self.register_buffer(name, val.to(torch.float32)) + + register_buffer('betas', betas) + register_buffer('alphas_cumprod', alphas_cumprod) + register_buffer('alphas_cumprod_prev', alphas_cumprod_prev) # calculations for diffusion q(x_t | x_{t-1}) and others - self.register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod)) - self.register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1. - alphas_cumprod)) - self.register_buffer('log_one_minus_alphas_cumprod', torch.log(1. - alphas_cumprod)) - self.register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. / alphas_cumprod)) - self.register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1. / alphas_cumprod - 1)) + register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod)) + register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1. - alphas_cumprod)) + register_buffer('log_one_minus_alphas_cumprod', torch.log(1. - alphas_cumprod)) + register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. / alphas_cumprod)) + register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1. / alphas_cumprod - 1)) # calculations for posterior q(x_{t-1} | x_t, x_0) @@ -365,13 +369,13 @@ class GaussianDiffusion(nn.Module): # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t) - self.register_buffer('posterior_variance', posterior_variance) + register_buffer('posterior_variance', posterior_variance) # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain - self.register_buffer('posterior_log_variance_clipped', torch.log(posterior_variance.clamp(min =1e-20))) - self.register_buffer('posterior_mean_coef1', betas * torch.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)) - self.register_buffer('posterior_mean_coef2', (1. - alphas_cumprod_prev) * torch.sqrt(alphas) / (1. - alphas_cumprod)) + register_buffer('posterior_log_variance_clipped', torch.log(posterior_variance.clamp(min =1e-20))) + register_buffer('posterior_mean_coef1', betas * torch.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)) + register_buffer('posterior_mean_coef2', (1. - alphas_cumprod_prev) * torch.sqrt(alphas) / (1. - alphas_cumprod)) def q_mean_variance(self, x_start, t): mean = extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start diff --git a/setup.py b/setup.py index 7659554..2c2d068 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.12.0', + version = '0.12.1', license='MIT', description = 'Denoising Diffusion Probabilistic Models - Pytorch', author = 'Phil Wang',