From ee76497b832d7d7e208dc2850ce614f08fe4a959 Mon Sep 17 00:00:00 2001 From: Phil Wang Date: Sat, 5 Sep 2020 23:41:02 -0700 Subject: [PATCH] complete training, sans sampling --- denoising_diffusion_pytorch/__init__.py | 2 +- .../denoising_diffusion_pytorch.py | 82 ++++++++++++++++++- 2 files changed, 80 insertions(+), 4 deletions(-) diff --git a/denoising_diffusion_pytorch/__init__.py b/denoising_diffusion_pytorch/__init__.py index 851bf72..f78d79f 100644 --- a/denoising_diffusion_pytorch/__init__.py +++ b/denoising_diffusion_pytorch/__init__.py @@ -1 +1 @@ -from denoising_diffusion_pytorch.denoising_diffusion_pytorch import DenoisingDiffusion, Unet +from denoising_diffusion_pytorch.denoising_diffusion_pytorch import GaussianDiffusion, Unet diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py index 23ece91..4b33308 100644 --- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py +++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py @@ -1,5 +1,7 @@ import math import torch +from inspect import isfunction +from functools import partial from torch import nn, einsum import torch.nn.functional as F @@ -12,7 +14,12 @@ def exists(x): return x is not None def default(val, d): - return d if not exists(val) else val + if exists(val): + return val + return d() if isfunction(d) else d + +def normal_kl(mean1, logvar1, mean2, logvar2): + return 0.5 * (-1. + logvar2 - logvar1 + torch.exp(logvar1 - logvar2) + torch.exp(-logvar2) * (mean1 - mean2) ** 2) # small helper modules @@ -30,9 +37,10 @@ class SinusoidalPosEmb(nn.Module): self.dim = dim def forward(self, x): + device = x.device half_dim = self.dim // 2 emb = math.log(10000) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim) * -emb) + emb = torch.exp(torch.arange(half_dim, device=device) * -emb) emb = x[:, None] * emb[None, :] emb = torch.cat((emb.sin(), emb.cos()), dim=-1) return emb @@ -189,7 +197,75 @@ class Unet(nn.Module): # gaussian diffusion trainer class -class DenoisingDiffusion(nn.Module): +def extract(a, t, x_shape): + b, *_ = t.shape + out = a.gather(-1, t) + return out.reshape(b, *((1,) * (len(x_shape) - 1))) + +class GaussianDiffusion(nn.Module): + def __init__(self, beta_start=0.0001, beta_end=0.02, num_diffusion_timesteps=1000, loss_type='l1'): + super().__init__() + self.np_betas = betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps).astype(np.float64) + timesteps, = betas.shape + self.num_timesteps = int(timesteps) + self.loss_type = loss_type + + alphas = 1. - betas + alphas_cumprod = np.cumprod(alphas, axis=0) + alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1]) + + to_torch = partial(torch.tensor, dtype=torch.float32) + + self.register_buffer('betas', to_torch(betas)) + self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod)) + self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev)) + + # calculations for diffusion q(x_t | x_{t-1}) and others + self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod))) + self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod))) + self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod))) + self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod))) + self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1))) + + # calculations for posterior q(x_{t-1} | x_t, x_0) + posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod) + # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t) + self.register_buffer('posterior_variance', to_torch(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', to_torch(np.log(np.maximum(posterior_variance, 1e-20)))) + self.register_buffer('posterior_mean_coef1', to_torch( + betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod))) + self.register_buffer('posterior_mean_coef2', to_torch( + (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod))) + + def q_sample(self, x_start, t, noise=None): + noise = default(noise, lambda: torch.randn_like(x_start)) + + return ( + extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + + extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise + ) + + def p_losses(self, x_start, t, denoise_fn, noise = None): + b, c, h, w = x_start.shape + noise = default(noise, lambda: torch.randn_like(x_start)) + + x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) + x_recon = denoise_fn(x_noisy, t) + + if self.loss_type == 'l1': + loss = (noise - x_recon).abs().mean() + elif self.loss_type == 'l2': + loss = F.mse_loss(noise, x_recon) + else: + raise NotImplementedError() + + return loss + + def forward(self, *args, **kwargs): + return self.p_losses(*args, **kwargs) + +class Trainer(nn.Module): def __init__(self): super().__init__()