optimize for simplicity and clarity - researcher does not need to worry about normalizing and unnormalizing now

denoising_diffusion_pytorch/__init__.py

@@ -1 +1,4 @@
 from denoising_diffusion_pytorch.denoising_diffusion_pytorch import GaussianDiffusion, Unet, Trainer
+
+from denoising_diffusion_pytorch.learned_gaussian_diffusion import LearnedGaussianDiffusion
+from denoising_diffusion_pytorch.weighted_objective_gaussian_diffusion import WeightedObjectiveGaussianDiffusion

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -40,6 +40,12 @@ def num_to_groups(num, divisor):
         arr.append(remainder)
     return arr
 
+def normalize_to_neg_one_to_one(img):
+    return img * 2 - 1
+
+def unnormalize_to_zero_to_one(t):
+    return (t + 1) * 0.5
+
 # small helper modules
 
 class EMA():
@@ -204,7 +210,8 @@ class Unet(nn.Module):
         dim_mults=(1, 2, 4, 8),
         channels = 3,
         with_time_emb = True,
-        resnet_block_groups = 8
+        resnet_block_groups = 8,
+        learned_variance = False
     ):
         super().__init__()
 
@@ -265,10 +272,12 @@ class Unet(nn.Module):
                 Upsample(dim_in) if not is_last else nn.Identity()
             ]))
 
-        out_dim = default(out_dim, channels)
+        default_out_dim = channels * (1 if not learned_variance else 2)
+        self.out_dim = default(out_dim, default_out_dim)
+
         self.final_conv = nn.Sequential(
             block_klass(dim, dim),
-            nn.Conv2d(dim, out_dim, 1)
+            nn.Conv2d(dim, self.out_dim, 1)
         )
 
     def forward(self, x, time):
@@ -320,7 +329,7 @@ def cosine_beta_schedule(timesteps, s = 0.008):
     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.9999)
+    return torch.clip(betas, 0, 0.999)
 
 class GaussianDiffusion(nn.Module):
     def __init__(
@@ -330,12 +339,16 @@ class GaussianDiffusion(nn.Module):
         image_size,
         channels = 3,
         timesteps = 1000,
-        loss_type = 'l1'
+        loss_type = 'l1',
+        objective = 'pred_noise'
     ):
         super().__init__()
+        assert not (type(self) == GaussianDiffusion and denoise_fn.channels != denoise_fn.out_dim)
+
         self.channels = channels
         self.image_size = image_size
         self.denoise_fn = denoise_fn
+        self.objective = objective
 
         betas = cosine_beta_schedule(timesteps)
 
@@ -377,12 +390,6 @@ class GaussianDiffusion(nn.Module):
         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
-        variance = extract(1. - self.alphas_cumprod, t, x_start.shape)
-        log_variance = extract(self.log_one_minus_alphas_cumprod, t, x_start.shape)
-        return mean, variance, log_variance
-
     def predict_start_from_noise(self, x_t, t, noise):
         return (
             extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
@@ -399,12 +406,19 @@ class GaussianDiffusion(nn.Module):
         return posterior_mean, posterior_variance, posterior_log_variance_clipped
 
     def p_mean_variance(self, x, t, clip_denoised: bool):
-        x_recon = self.predict_start_from_noise(x, t=t, noise=self.denoise_fn(x, t))
+        model_output = self.denoise_fn(x, t)
+
+        if self.objective == 'pred_noise':
+            x_start = self.predict_start_from_noise(x, t = t, noise = model_output)
+        elif self.objective == 'pred_x0':
+            x_start = model_output
+        else:
+            raise ValueError(f'unknown objective {self.objective}')
 
         if clip_denoised:
-            x_recon.clamp_(-1., 1.)
+            x_start.clamp_(-1., 1.)
 
-        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
+        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start = x_start, x_t = x, t = t)
         return model_mean, posterior_variance, posterior_log_variance
 
     @torch.no_grad()
@@ -425,6 +439,8 @@ class GaussianDiffusion(nn.Module):
 
         for i in tqdm(reversed(range(0, self.num_timesteps)), desc='sampling loop time step', total=self.num_timesteps):
             img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long))
+
+        img = normalize_to_neg_one_to_one(img)
         return img
 
     @torch.no_grad()
@@ -457,27 +473,39 @@ class GaussianDiffusion(nn.Module):
             extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
         )
 
+    @property
+    def loss_fn(self):
+        if self.loss_type == 'l1':
+            return F.l1_loss
+        elif self.loss_type == 'l2':
+            return F.mse_loss
+        else:
+            raise ValueError(f'invalid loss type {self.loss_type}')
+
     def p_losses(self, x_start, t, 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 = self.denoise_fn(x_noisy, t)
+        x = self.q_sample(x_start=x_start, t=t, noise=noise)
+        model_out = self.denoise_fn(x, t)
 
-        if self.loss_type == 'l1':
-            loss = (noise - x_recon).abs().mean()
-        elif self.loss_type == 'l2':
-            loss = F.mse_loss(noise, x_recon)
+        if self.objective == 'pred_noise':
+            target = noise
+        elif self.objective == 'pred_x0':
+            target = x_start
         else:
-            raise NotImplementedError()
+            raise ValueError(f'unknown objective {self.objective}')
 
+        loss = self.loss_fn(model_out, target)
         return loss
 
-    def forward(self, x, *args, **kwargs):
-        b, c, h, w, device, img_size, = *x.shape, x.device, self.image_size
+    def forward(self, img, *args, **kwargs):
+        b, c, h, w, device, img_size, = *img.shape, img.device, self.image_size
         assert h == img_size and w == img_size, f'height and width of image must be {img_size}'
         t = torch.randint(0, self.num_timesteps, (b,), device=device).long()
-        return self.p_losses(x, t, *args, **kwargs)
+
+        img = normalize_to_neg_one_to_one(img)
+        return self.p_losses(img, t, *args, **kwargs)
 
 # dataset classes
 
@@ -492,8 +520,7 @@ class Dataset(data.Dataset):
             transforms.Resize(image_size),
             transforms.RandomHorizontalFlip(),
             transforms.CenterCrop(image_size),
-            transforms.ToTensor(),
-            transforms.Lambda(lambda t: (t * 2) - 1)
+            transforms.ToTensor()
         ])
 
     def __len__(self):
@@ -579,32 +606,36 @@ class Trainer(object):
         self.scaler.load_state_dict(data['scaler'])
 
     def train(self):
-        while self.step < self.train_num_steps:
-            for i in range(self.gradient_accumulate_every):
-                data = next(self.dl).cuda()
+        with tqdm(initial = self.step, total = self.train_num_steps) as pbar:
 
-                with autocast(enabled = self.amp):
-                    loss = self.model(data)
-                    self.scaler.scale(loss / self.gradient_accumulate_every).backward()
+            while self.step < self.train_num_steps:
+                for i in range(self.gradient_accumulate_every):
+                    data = next(self.dl).cuda()
 
-                print(f'{self.step}: {loss.item()}')
+                    with autocast(enabled = self.amp):
+                        loss = self.model(data)
+                        self.scaler.scale(loss / self.gradient_accumulate_every).backward()
 
-            self.scaler.step(self.opt)
-            self.scaler.update()
-            self.opt.zero_grad()
+                    pbar.set_description(f'loss: {loss.item():.4f}')
 
-            if self.step % self.update_ema_every == 0:
-                self.step_ema()
+                self.scaler.step(self.opt)
+                self.scaler.update()
+                self.opt.zero_grad()
 
-            if self.step != 0 and self.step % self.save_and_sample_every == 0:
-                milestone = self.step // self.save_and_sample_every
-                batches = num_to_groups(36, self.batch_size)
-                all_images_list = list(map(lambda n: self.ema_model.sample(batch_size=n), batches))
-                all_images = torch.cat(all_images_list, dim=0)
-                all_images = (all_images + 1) * 0.5
-                utils.save_image(all_images, str(self.results_folder / f'sample-{milestone}.png'), nrow = 6)
-                self.save(milestone)
+                if self.step % self.update_ema_every == 0:
+                    self.step_ema()
 
-            self.step += 1
+                if self.step != 0 and self.step % self.save_and_sample_every == 0:
+                    self.ema_model.eval()
 
-        print('training completed')
+                    milestone = self.step // self.save_and_sample_every
+                    batches = num_to_groups(36, self.batch_size)
+                    all_images_list = list(map(lambda n: self.ema_model.sample(batch_size=n), batches))
+                    all_images = torch.cat(all_images_list, dim=0)
+                    utils.save_image(all_images, str(self.results_folder / f'sample-{milestone}.png'), nrow = 6)
+                    self.save(milestone)
+
+                self.step += 1
+                pbar.update(1)
+
+        print('training complete')

denoising_diffusion_pytorch/learned_gaussian_diffusion.py

@@ -0,0 +1,132 @@
+import torch
+from math import pi, sqrt, log as ln
+from inspect import isfunction
+from torch import nn, einsum
+from einops import rearrange
+
+from denoising_diffusion_pytorch.denoising_diffusion_pytorch import GaussianDiffusion, extract, unnormalize_to_zero_to_one
+
+# constants
+
+NAT = 1. / ln(2)
+
+# helper functions
+
+def exists(x):
+    return x is not None
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+# tensor helpers
+
+def log(t, eps = 1e-12):
+    return torch.log(t.clamp(min = eps))
+
+def meanflat(x):
+    return x.mean(dim = tuple(range(1, len(x.shape))))
+
+def normal_kl(mean1, logvar1, mean2, logvar2):
+    """
+    KL divergence between normal distributions parameterized by mean and log-variance.
+    """
+    return 0.5 * (-1.0 + logvar2 - logvar1 + torch.exp(logvar1 - logvar2) + ((mean1 - mean2) ** 2) * torch.exp(-logvar2))
+
+def approx_standard_normal_cdf(x):
+    return 0.5 * (1.0 + torch.tanh(sqrt(2.0 / pi) * (x + 0.044715 * (x ** 3))))
+
+def discretized_gaussian_log_likelihood(x, *, means, log_scales, thres = 0.999):
+    assert x.shape == means.shape == log_scales.shape
+
+    centered_x = x - means
+    inv_stdv = torch.exp(-log_scales)
+    plus_in = inv_stdv * (centered_x + 1. / 255.)
+    cdf_plus = approx_standard_normal_cdf(plus_in)
+    min_in = inv_stdv * (centered_x - 1. / 255.)
+    cdf_min = approx_standard_normal_cdf(min_in)
+    log_cdf_plus = log(cdf_plus)
+    log_one_minus_cdf_min = log(1. - cdf_min)
+    cdf_delta = cdf_plus - cdf_min
+
+    log_probs = torch.where(x < -thres,
+        log_cdf_plus,
+        torch.where(x > thres,
+            log_one_minus_cdf_min,
+            log(cdf_delta)))
+
+    return log_probs
+
+# https://arxiv.org/abs/2102.09672
+
+# i thought the results were questionable, if one were to focus only on FID
+# but may as well get this in here for others to try, as GLIDE is using it (and DALL-E2 first stage of cascade)
+# gaussian diffusion for learned variance + hybrid eps simple + vb loss
+
+class LearnedGaussianDiffusion(GaussianDiffusion):
+    def __init__(
+        self,
+        denoise_fn,
+        vb_loss_weight = 0.001,  # lambda was 0.001 in the paper
+        *args,
+        **kwargs
+    ):
+        super().__init__(denoise_fn, *args, **kwargs)
+        assert denoise_fn.out_dim == (denoise_fn.channels * 2), 'dimension out of unet must be twice the number of channels for learned variance - you can also set the `learned_variance` keyword argument on the Unet to be `True`'
+        self.vb_loss_weight = vb_loss_weight
+
+    def p_mean_variance(self, *, x, t, clip_denoised, model_output = None):
+        model_output = default(model_output, lambda: self.denoise_fn(x, t))
+        pred_noise, var_interp_frac_unnormalized = model_output.chunk(2, dim = 1)
+
+        min_log = extract(self.posterior_log_variance_clipped, t, x.shape)
+        max_log = extract(torch.log(self.betas), t, x.shape)
+        var_interp_frac = unnormalize_to_zero_to_one(var_interp_frac_unnormalized)
+
+        model_log_variance = var_interp_frac * max_log + (1 - var_interp_frac) * min_log
+        model_variance = model_log_variance.exp()
+
+        x_start = self.predict_start_from_noise(x, t, pred_noise)
+
+        if clip_denoised:
+            x_start.clamp_(-1., 1.)
+
+        model_mean, _, _ = self.q_posterior(x_start, x, t)
+
+        return model_mean, model_variance, model_log_variance
+
+    def p_losses(self, x_start, t, noise = None, clip_denoised = False):
+        noise = default(noise, lambda: torch.randn_like(x_start))
+        x_t = self.q_sample(x_start = x_start, t = t, noise = noise)
+
+        # model output
+
+        model_output = self.denoise_fn(x_t, t)
+
+        # calculating kl loss for learned variance (interpolation)
+
+        true_mean, _, true_log_variance_clipped = self.q_posterior(x_start = x_start, x_t = x_t, t = t)
+        model_mean, _, model_log_variance = self.p_mean_variance(x = x_t, t = t, clip_denoised = clip_denoised, model_output = model_output)
+
+        # kl loss with detached model predicted mean, for stability reasons as in paper
+
+        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 = 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))
+
+        vb_losses = torch.where(t == 0, decoder_nll, kl)
+
+        # simple loss - predicting noise, x0, or x_prev
+
+        pred_noise, _ = model_output.chunk(2, dim = 1)
+
+        simple_losses = self.loss_fn(pred_noise, noise)
+
+        return simple_losses + vb_losses.mean() * self.vb_loss_weight

denoising_diffusion_pytorch/weighted_objective_gaussian_diffusion.py

@@ -0,0 +1,80 @@
+import torch
+from inspect import isfunction
+from torch import nn, einsum
+from einops import rearrange
+
+from denoising_diffusion_pytorch.denoising_diffusion_pytorch import GaussianDiffusion
+
+# helper functions
+
+def exists(x):
+    return x is not None
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+# some improvisation on my end
+# where i have the model learn to both predict noise and x0
+# and learn the weighted sum for each depending on time step
+
+class WeightedObjectiveGaussianDiffusion(GaussianDiffusion):
+    def __init__(
+        self,
+        denoise_fn,
+        *args,
+        pred_noise_loss_weight = 0.1,
+        pred_x_start_loss_weight = 0.1,
+        **kwargs
+    ):
+        super().__init__(denoise_fn, *args, **kwargs)
+        channels = denoise_fn.channels
+        assert denoise_fn.out_dim == (channels * 2 + 2), 'dimension out (out_dim) of unet must be twice the number of channels + 2 (for the softmax weighted sum) - for channels of 3, this should be (3 * 2) + 2 = 8'
+
+        self.split_dims = (channels, channels, 2)
+        self.pred_noise_loss_weight = pred_noise_loss_weight
+        self.pred_x_start_loss_weight = pred_x_start_loss_weight
+
+    def p_mean_variance(self, *, x, t, clip_denoised, model_output = None):
+        model_output = self.denoise_fn(x, t)
+
+        pred_noise, pred_x_start, weights = model_output.split(self.split_dims, dim = 1)
+        normalized_weights = weights.softmax(dim = 1)
+
+        x_start_from_noise = self.predict_start_from_noise(x, t = t, noise = pred_noise)
+        
+        x_starts = torch.stack((x_start_from_noise, pred_x_start), dim = 1)
+        weighted_x_start = einsum('b j h w, b j c h w -> b c h w', normalized_weights, x_starts)
+
+        if clip_denoised:
+            weighted_x_start.clamp_(-1., 1.)
+
+        model_mean, model_variance, model_log_variance = self.q_posterior(weighted_x_start, x, t)
+
+        return model_mean, model_variance, model_log_variance
+
+    def p_losses(self, x_start, t, noise = None, clip_denoised = False):
+        noise = default(noise, lambda: torch.randn_like(x_start))
+        x_t = self.q_sample(x_start = x_start, t = t, noise = noise)
+
+        model_output = self.denoise_fn(x_t, t)
+        pred_noise, pred_x_start, weights = model_output.split(self.split_dims, dim = 1)
+
+        # get loss for predicted noise and x_start
+        # with the loss weight given at initialization
+
+        noise_loss = self.loss_fn(noise, pred_noise) * self.pred_noise_loss_weight
+        x_start_loss = self.loss_fn(x_start, pred_x_start) * self.pred_x_start_loss_weight
+
+        # calculate x_start from predicted noise
+        # then do a weighted sum of the x_start prediction, weights also predicted by the model (softmax normalized)
+
+        x_start_from_pred_noise = self.predict_start_from_noise(x_t, t, pred_noise)
+        x_start_from_pred_noise = x_start_from_pred_noise.clamp(-2., 2.)
+        weighted_x_start = einsum('b j h w, b j c h w -> b c h w', weights.softmax(dim = 1), torch.stack((x_start_from_pred_noise, pred_x_start), dim = 1))
+
+        # main loss to x_start with the weighted one
+
+        weighted_x_start_loss = self.loss_fn(x_start, weighted_x_start)
+        return weighted_x_start_loss + x_start_loss + noise_loss

setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'denoising-diffusion-pytorch',
   packages = find_packages(),
-  version = '0.12.1',
+  version = '0.15.4',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',