save samples and models to ./results path

README.md

@@ -6,6 +6,8 @@ Implementation of <a href="https://arxiv.org/abs/2006.11239">Denoising Diffusion
 
 <img src="./sample.png" width="500px"><img>
 
+[![PyPI version](https://badge.fury.io/py/denoising-diffusion-pytorch.svg)](https://badge.fury.io/py/denoising-diffusion-pytorch)
+
 ## Install
 
 ```bash
@@ -25,10 +27,8 @@ model = Unet(
 
 diffusion = GaussianDiffusion(
     model,
-    beta_start = 0.0001,
-    beta_end = 0.02,
-    num_diffusion_timesteps = 1000,   # number of steps
-    loss_type = 'l1'                  # L1 or L2 (wavegrad paper claims l1 is better?)
+    timesteps = 1000,   # number of steps
+    loss_type = 'l1'    # L1 or L2
 )
 
 training_images = torch.randn(8, 3, 128, 128)
@@ -52,10 +52,8 @@ model = Unet(
 
 diffusion = GaussianDiffusion(
     model,
-    beta_start = 0.0001,
-    beta_end = 0.02,
-    num_diffusion_timesteps = 1000,   # number of steps
-    loss_type = 'l1'                  # L1 or L2
+    timesteps = 1000,   # number of steps
+    loss_type = 'l1'    # L1 or L2
 ).cuda()
 
 trainer = Trainer(
@@ -66,14 +64,13 @@ trainer = Trainer(
     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
+    ema_decay = 0.995,                # exponential moving average decay
+    fp16 = True                       # turn on mixed precision training with apex
 )
 
 trainer.train()
 ```
 
-Todo: Command line tool for one-line training
-
 ## Citations
 
 ```bibtex
@@ -86,3 +83,15 @@ Todo: Command line tool for one-line training
     primaryClass={cs.LG}
 }
 ```
+
+```bibtex
+@inproceedings{
+    anonymous2021improved,
+    title={Improved Denoising Diffusion Probabilistic Models},
+    author={Anonymous},
+    booktitle={Submitted to International Conference on Learning Representations},
+    year={2021},
+    url={https://openreview.net/forum?id=-NEXDKk8gZ},
+    note={under review}
+}
+```

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -16,11 +16,20 @@ import numpy as np
 from tqdm import tqdm
 from einops import rearrange
 
+try:
+    from apex import amp
+    APEX_AVAILABLE = True
+except:
+    APEX_AVAILABLE = False
+
 # constants
 
 SAVE_AND_SAMPLE_EVERY = 1000
 UPDATE_EMA_EVERY = 10
-EXTS = ['jpg', 'png']
+EXTS = ['jpg', 'jpeg', 'png']
+
+RESULTS_FOLDER = Path('./results')
+RESULTS_FOLDER.mkdir(exist_ok = True)
 
 # helpers functions
 
@@ -37,6 +46,21 @@ def cycle(dl):
         for data in dl:
             yield data
 
+def num_to_groups(num, divisor):
+    groups = num // divisor
+    remainder = num % divisor
+    arr = [divisor] * groups
+    if remainder > 0:
+        arr.append(remainder)
+    return arr
+
+def loss_backwards(fp16, loss, optimizer, **kwargs):
+    if fp16:
+        with amp.scale_loss(loss, optimizer) as scaled_loss:
+            scaled_loss.backward(**kwargs)
+    else:
+        loss.backward(**kwargs)
+
 # small helper modules
 
 class EMA():
@@ -97,17 +121,18 @@ class Downsample(nn.Module):
         return self.conv(x)
 
 class Rezero(nn.Module):
-    def __init__(self, dim):
+    def __init__(self, fn):
         super().__init__()
+        self.fn = fn
         self.g = nn.Parameter(torch.zeros(1))
 
     def forward(self, x):
-        return x * self.g
+        return self.fn(x) * self.g
 
 # building block modules
 
 class Block(nn.Module):
-    def __init__(self, dim, dim_out, groups = 32):
+    def __init__(self, dim, dim_out, groups = 8):
         super().__init__()
         self.block = nn.Sequential(
             nn.Conv2d(dim, dim_out, 3, padding=1),
@@ -118,7 +143,7 @@ class Block(nn.Module):
         return self.block(x)
 
 class ResnetBlock(nn.Module):
-    def __init__(self, dim, dim_out, *, time_emb_dim, groups = 32):
+    def __init__(self, dim, dim_out, *, time_emb_dim, groups = 8):
         super().__init__()
         self.mlp = nn.Sequential(
             Mish(),
@@ -136,18 +161,17 @@ class ResnetBlock(nn.Module):
         return h + self.res_conv(x)
 
 class LinearAttention(nn.Module):
-    def __init__(self, dim, heads = 8, dim_head = 32):
+    def __init__(self, dim, heads = 4, dim_head = 32):
         super().__init__()
         self.heads = heads
         hidden_dim = dim_head * heads
-        self.to_qkv = nn.Conv2d(dim, hidden_dim, 1, bias = False)
+        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False)
         self.to_out = nn.Conv2d(hidden_dim, dim, 1)
 
     def forward(self, x):
         b, c, h, w = x.shape
         qkv = self.to_qkv(x)
-        q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads)
-        q = q.softmax(dim=-2)
+        q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads, qkv=3)
         k = k.softmax(dim=-1)
         context = torch.einsum('bhdn,bhen->bhde', k, v)
         out = torch.einsum('bhde,bhdn->bhen', context, q)
@@ -157,7 +181,7 @@ class LinearAttention(nn.Module):
 # model
 
 class Unet(nn.Module):
-    def __init__(self, dim, out_dim = None, dim_mults=(1, 2, 4, 8), groups = 32):
+    def __init__(self, dim, out_dim = None, dim_mults=(1, 2, 4, 8), groups = 8):
         super().__init__()
         dims = [3, *map(lambda m: dim * m, dim_mults)]
         in_out = list(zip(dims[:-1], dims[1:]))
@@ -178,6 +202,7 @@ class Unet(nn.Module):
 
             self.downs.append(nn.ModuleList([
                 ResnetBlock(dim_in, dim_out, time_emb_dim = dim),
+                ResnetBlock(dim_out, dim_out, time_emb_dim = dim),
                 Residual(Rezero(LinearAttention(dim_out))),
                 Downsample(dim_out) if not is_last else nn.Identity()
             ]))
@@ -192,6 +217,7 @@ class Unet(nn.Module):
 
             self.ups.append(nn.ModuleList([
                 ResnetBlock(dim_out * 2, dim_in, time_emb_dim = dim),
+                ResnetBlock(dim_in, dim_in, time_emb_dim = dim),
                 Residual(Rezero(LinearAttention(dim_in))),
                 Upsample(dim_in) if not is_last else nn.Identity()
             ]))
@@ -208,8 +234,9 @@ class Unet(nn.Module):
 
         h = []
 
-        for resnet, attn, downsample in self.downs:
+        for resnet, resnet2, attn, downsample in self.downs:
             x = resnet(x, t)
+            x = resnet2(x, t)
             x = attn(x)
             h.append(x)
             x = downsample(x)
@@ -218,9 +245,10 @@ class Unet(nn.Module):
         x = self.mid_attn(x)
         x = self.mid_block2(x, t)
 
-        for resnet, attn, upsample in self.ups:
+        for resnet, resnet2, attn, upsample in self.ups:
             x = torch.cat((x, h.pop()), dim=1)
             x = resnet(x, t)
+            x = resnet2(x, t)
             x = attn(x)
             x = upsample(x)
 
@@ -238,24 +266,36 @@ def noise_like(shape, device, repeat=False):
     noise = lambda: torch.randn(shape, device=device)
     return repeat_noise() if repeat else noise()
 
+def cosine_beta_schedule(timesteps, s = 0.008):
+    """
+    cosine schedule
+    as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
+    """
+    steps = timesteps + 1
+    x = np.linspace(0, steps, steps)
+    alphas_cumprod = np.cos(((x / steps) + s) / (1 + s) * np.pi * 0.5) ** 2
+    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
+    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
+    return np.clip(betas, a_min = 0, a_max = 0.999)
+
 class GaussianDiffusion(nn.Module):
-    def __init__(self, denoise_fn, beta_start=0.0001, beta_end=0.02, num_diffusion_timesteps=1000, loss_type='l1', betas = None):
+    def __init__(self, denoise_fn, timesteps=1000, loss_type='l1', betas = None):
         super().__init__()
         self.denoise_fn = denoise_fn
 
         if exists(betas):
-            self.np_betas = betas.detach().cpu().numpy() if isinstance(betas, torch.Tensor) else betas
+            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
+            betas = cosine_beta_schedule(timesteps)
 
         alphas = 1. - betas
         alphas_cumprod = np.cumprod(alphas, axis=0)
         alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
 
+        timesteps, = betas.shape
+        self.num_timesteps = int(timesteps)
+        self.loss_type = loss_type
+
         to_torch = partial(torch.tensor, dtype=torch.float32)
 
         self.register_buffer('betas', to_torch(betas))
@@ -332,7 +372,7 @@ class GaussianDiffusion(nn.Module):
 
     @torch.no_grad()
     def sample(self, image_size, batch_size = 16):
-        return self.p_sample_loop((16, 3, image_size, image_size))
+        return self.p_sample_loop((batch_size, 3, image_size, image_size))
 
     @torch.no_grad()
     def interpolate(self, x1, x2, t = None, lam = 0.5):
@@ -417,55 +457,80 @@ class Trainer(object):
         train_lr = 2e-5,
         train_num_steps = 100000,
         gradient_accumulate_every = 2,
+        fp16 = False,
+        step_start_ema = 2000
     ):
         super().__init__()
         self.model = diffusion_model
+        self.ema = EMA(ema_decay)
+        self.ema_model = copy.deepcopy(self.model)
+        self.step_start_ema = step_start_ema
 
+        self.batch_size = train_batch_size
         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
 
+        assert not fp16 or fp16 and APEX_AVAILABLE, 'Apex must be installed in order for mixed precision training to be turned on'
+
+        self.fp16 = fp16
+        if fp16:
+            (self.model, self.ema_model), self.opt = amp.initialize([self.model, self.ema_model], self.opt, opt_level='O1')
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        self.ema_model.load_state_dict(self.model.state_dict())
+
+    def step_ema(self):
+        if self.step < self.step_start_ema:
+            self.reset_parameters()
+            return
+        self.ema.update_model_average(self.ema_model, self.model)
+
     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')
+        torch.save(data, str(RESULTS_FOLDER / f'model-{milestone}.pt'))
 
     def load(self, milestone):
-        data = torch.load(f'./model-{milestone}.pt')
+        data = torch.load(str(RESULTS_FOLDER / 'model-{milestone}.pt'))
+
         self.step = data['step']
-        self.model = data['model']
-        self.ema_model = data['ema']
+        self.model.load_state_dict(data['model'])
+        self.ema_model.load_state_dict(data['ema'])
 
     def train(self):
+        backwards = partial(loss_backwards, self.fp16)
+
         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()
+                backwards(loss / self.gradient_accumulate_every, self.opt)
 
             self.opt.step()
             self.opt.zero_grad()
 
             if self.step % UPDATE_EMA_EVERY == 0:
-                self.ema.update_model_average(self.ema_model, self.model)
+                self.step_ema()
 
-            if self.step % SAVE_AND_SAMPLE_EVERY == 0:
+            if self.step != 0 and 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)
+                batches = num_to_groups(36, self.batch_size)
+                all_images_list = list(map(lambda n: self.ema_model.sample(self.image_size, batch_size=n), batches))
+                all_images = torch.cat(all_images_list, dim=0)
+                utils.save_image(all_images, str(RESULTS_FOLDER / 'sample-{milestone}.png'), nrow=6)
                 self.save(milestone)
 
             self.step += 1

setup.py

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