add exponential moving average of model, also allow passing in custom noise schedule

README.md

@@ -1,6 +1,10 @@
-## Denoising Diffusion Probabilistic Model, in Pytorch (wip)
+<img src="./denoising-diffusion.png" width="500px"></img>
 
-Implementation of <a href="https://arxiv.org/abs/2006.11239">Denoising Diffusion Probabilistic Model</a> in Pytorch.
+## Denoising Diffusion Probabilistic Model, in Pytorch
+
+Implementation of <a href="https://arxiv.org/abs/2006.11239">Denoising Diffusion Probabilistic Model</a> in Pytorch. It is a new approach to generative modeling that may <a href="https://ajolicoeur.wordpress.com/the-new-contender-to-gans-score-matching-with-langevin-sampling/">have the potential</a> to rival GANs. It uses denoising score matching to estimate the gradient of the data distribution, followed by Langevin sampling to sample from the true distribution. This implementation was transcribed from the official Tensorflow version <a href="https://github.com/hojonathanho/diffusion">here</a>. 
+
+<img src="./sample.png" width="500px"><img>
 
 ## Install
 
@@ -32,10 +36,44 @@ loss = diffusion(training_images)
 loss.backward()
 # after a lot of training
 
-sampled_images = diffusion.p_sample_loop((1, 3, 128, 128))
-sampled_images.shape # (1, 3, 128, 128)
+sampled_images = diffusion.sample(128, batch_size = 4)
+sampled_images.shape # (4, 3, 128, 128)
 ```
 
+Or, if you simply want to pass in a folder name and the desired image dimensions, you can use the `Trainer` class to easily train a model.
+
+```python
+from denoising_diffusion_pytorch import Unet, GaussianDiffusion, Trainer
+
+model = Unet(
+    dim = 64,
+    dim_mults = (1, 2, 4, 8)
+).cuda()
+
+diffusion = GaussianDiffusion(
+    model,
+    beta_start = 0.0001,
+    beta_end = 0.02,
+    num_diffusion_timesteps = 1000,   # number of steps
+    loss_type = 'l1'                  # L1 or L2
+).cuda()
+
+trainer = Trainer(
+    diffusion,
+    'path/to/your/images',
+    image_size = 128,
+    train_batch_size = 32,
+    train_lr = 2e-5,
+    train_num_steps = 100000,         # total training steps
+    gradient_accumulate_every = 2,    # gradient accumulation steps
+    ema_decay = 0.995                 # exponential moving average decay
+)
+
+trainer.train()
+```
+
+Todo: Command line tool for one-line training
+
 ## Citations
 
 ```bibtex
@@ -48,14 +86,3 @@ sampled_images.shape # (1, 3, 128, 128)
     primaryClass={cs.LG}
 }
 ```
-
-```bibtex
-@misc{chen2020wavegrad,
-    title={WaveGrad: Estimating Gradients for Waveform Generation},
-    author={Nanxin Chen and Yu Zhang and Heiga Zen and Ron J. Weiss and Mohammad Norouzi and William Chan},
-    year={2020},
-    eprint={2009.00713},
-    archivePrefix={arXiv},
-    primaryClass={eess.AS}
-}
-```

denoising_diffusion_pytorch/__init__.py

@@ -1 +1 @@
-from denoising_diffusion_pytorch.denoising_diffusion_pytorch import GaussianDiffusion, Unet
+from denoising_diffusion_pytorch.denoising_diffusion_pytorch import GaussianDiffusion, Unet, Trainer

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -1,14 +1,27 @@
 import math
+import copy
 import torch
-from inspect import isfunction
-from functools import partial
 from torch import nn, einsum
 import torch.nn.functional as F
+from inspect import isfunction
+from functools import partial
+
+from torch.utils import data
+from pathlib import Path
+from torch.optim import Adam
+from torchvision import transforms, utils
+from PIL import Image
 
 import numpy as np
 from tqdm import tqdm
 from einops import rearrange
 
+# constants
+
+SAVE_AND_SAMPLE_EVERY = 1000
+UPDATE_EMA_EVERY = 100
+EXTS = ['jpg', 'png']
+
 # helpers functions
 
 def exists(x):
@@ -19,11 +32,28 @@ def default(val, d):
         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)
+def cycle(dl):
+    while True:
+        for data in dl:
+            yield data
 
 # small helper modules
 
+class EMA():
+    def __init__(self, beta):
+        super().__init__()
+        self.beta = beta
+
+    def update_model_average(self, ma_model, current_model):
+        for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()):
+            old_weight, up_weight = ma_params.data, current_params.data
+            ma_params.data = self.update_average(old_weight, up_weight)
+
+    def update_average(self, old, new):
+        if old is None:
+            return new
+        return old * self.beta + (1 - self.beta) * new
+
 class Residual(nn.Module):
     def __init__(self, fn):
         super().__init__()
@@ -209,11 +239,15 @@ def noise_like(shape, device, repeat=False):
     return repeat_noise() if repeat else noise()
 
 class GaussianDiffusion(nn.Module):
-    def __init__(self, denoise_fn, beta_start=0.0001, beta_end=0.02, num_diffusion_timesteps=1000, loss_type='l1'):
+    def __init__(self, denoise_fn, beta_start=0.0001, beta_end=0.02, num_diffusion_timesteps=1000, loss_type='l1', betas = None):
         super().__init__()
         self.denoise_fn = denoise_fn
 
-        self.np_betas = betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps).astype(np.float64)
+        if exists(betas):
+            self.np_betas = betas.detach().cpu().numpy() if isinstance(betas, torch.Tensor) else betas
+        else:
+            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
@@ -296,6 +330,26 @@ class GaussianDiffusion(nn.Module):
             img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long))
         return img
 
+    @torch.no_grad()
+    def sample(self, image_size, batch_size = 16):
+        return self.p_sample_loop((16, 3, image_size, image_size))
+
+    @torch.no_grad()
+    def interpolate(self, x1, x2, t = None, lam = 0.5):
+        b, *_, device = *x1.shape, x1.device
+        t = default(t, self.num_timesteps - 1)
+
+        assert x1.shape == x2.shape
+
+        t_batched = torch.stack([torch.tensor(t, device=device)] * b)
+        xt1, xt2 = map(lambda x: self.q_sample(x, t=t_batched), (x1, x2))
+
+        img = (1 - lam) * xt1 + lam * xt2
+        for i in tqdm(reversed(range(0, t)), desc='interpolation sample time step', total=t):
+            img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long))
+
+        return img
+
     def q_sample(self, x_start, t, noise=None):
         noise = default(noise, lambda: torch.randn_like(x_start))
 
@@ -324,3 +378,96 @@ class GaussianDiffusion(nn.Module):
         b, *_, device = *x.shape, x.device
         t = torch.randint(0, self.num_timesteps, (b,), device=device).long()
         return self.p_losses(x, t, *args, **kwargs)
+
+# dataset classes
+
+class Dataset(data.Dataset):
+    def __init__(self, folder, image_size):
+        super().__init__()
+        self.folder = folder
+        self.image_size = image_size
+        self.paths = [p for ext in EXTS for p in Path(f'{folder}').glob(f'**/*.{ext}')]
+
+        self.transform = transforms.Compose([
+            transforms.Resize(image_size),
+            transforms.RandomHorizontalFlip(),
+            transforms.CenterCrop(image_size),
+            transforms.ToTensor()
+        ])
+
+    def __len__(self):
+        return len(self.paths)
+
+    def __getitem__(self, index):
+        path = self.paths[index]
+        img = Image.open(path)
+        return self.transform(img)
+
+# trainer class
+
+class Trainer(object):
+    def __init__(
+        self,
+        diffusion_model,
+        folder,
+        *,
+        ema_decay = 0.995,
+        image_size = 128,
+        train_batch_size = 32,
+        train_lr = 2e-5,
+        train_num_steps = 100000,
+        gradient_accumulate_every = 2,
+    ):
+        super().__init__()
+        self.model = diffusion_model
+
+        self.image_size = image_size
+        self.gradient_accumulate_every = gradient_accumulate_every
+        self.train_num_steps = train_num_steps
+
+        self.ema = EMA(ema_decay)
+        self.ema_model = copy.deepcopy(self.model)
+
+        self.ds = Dataset(folder, image_size)
+        self.dl = cycle(data.DataLoader(self.ds, batch_size = train_batch_size, shuffle=True, pin_memory=True))
+        self.opt = Adam(diffusion_model.parameters(), lr=train_lr)
+
+        self.step = 0
+
+    def save(self, milestone):
+        data = {
+            'step': self.step,
+            'model': self.model.state_dict(),
+            'ema': self.ema_model.state_dict()
+        }
+        torch.save(data, f'./model-{milestone}.pt')
+
+    def load(self, milestone):
+        data = torch.load(f'./model-{milestone}.pt')
+        self.step = data['step']
+        self.model = data['model']
+        self.ema_model = data['ema']
+
+    def train(self):
+        while self.step < self.train_num_steps:
+            for i in range(self.gradient_accumulate_every):
+                data = next(self.dl).cuda()
+                loss = self.model(data)
+                print(f'{self.step}: {loss.item()}')
+                (loss / self.gradient_accumulate_every).backward()
+
+            self.opt.step()
+            self.opt.zero_grad()
+
+            if self.step % UPDATE_EMA_EVERY == 0:
+                self.ema.update_model_average(self.ema_model, self.model)
+
+            if self.step % SAVE_AND_SAMPLE_EVERY == 0:
+                milestone = self.step // SAVE_AND_SAMPLE_EVERY
+                all_images = self.ema_model.p_sample_loop((64, 3, self.image_size, self.image_size))
+                utils.save_image(all_images, f'./sample-{milestone}.png', nrow=8)
+                self.save(milestone)
+
+            self.step += 1
+
+        print('training completed')

setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'denoising-diffusion-pytorch',
   packages = find_packages(),
-  version = '0.0.2',
+  version = '0.2.0',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',
@@ -16,7 +16,9 @@ setup(
   install_requires=[
     'einops',
     'numpy',
+    'pillow',
     'torch',
+    'torchvision',
     'tqdm'
   ],
   classifiers=[