save samples and models to ./results path

README.md

@@ -27,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)
@@ -54,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(
@@ -87,3 +83,15 @@ trainer.train()
     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

@@ -26,7 +26,10 @@ except:
 
 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
 
@@ -169,7 +172,6 @@ class LinearAttention(nn.Module):
         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, qkv=3)
-        q = q.softmax(dim=-2)
         k = k.softmax(dim=-1)
         context = torch.einsum('bhdn,bhen->bhde', k, v)
         out = torch.einsum('bhde,bhdn->bhen', context, q)
@@ -264,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))
@@ -486,10 +500,10 @@ class Trainer(object):
             '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 / f'model-{milestone}.pt'))
 
         self.step = data['step']
         self.model.load_state_dict(data['model'])
@@ -516,7 +530,7 @@ class Trainer(object):
                 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, f'./sample-{milestone}.png', nrow=6)
+                utils.save_image(all_images, str(RESULTS_FOLDER / f'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.3.2',
+  version = '0.5.2',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',