allow for mixed precision training with fp16 flag

README.md

@@ -6,8 +6,6 @@ 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
@@ -75,6 +73,8 @@ trainer = Trainer(
 trainer.train()
 ```
 
+Todo: Command line tool for one-line training
+
 ## Citations
 
 ```bibtex

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -43,14 +43,6 @@ 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:
@@ -118,13 +110,12 @@ class Downsample(nn.Module):
         return self.conv(x)
 
 class Rezero(nn.Module):
-    def __init__(self, fn):
+    def __init__(self, dim):
         super().__init__()
-        self.fn = fn
         self.g = nn.Parameter(torch.zeros(1))
 
     def forward(self, x):
-        return self.fn(x) * self.g
+        return x * self.g
 
 # building block modules
 
@@ -158,17 +149,17 @@ class ResnetBlock(nn.Module):
         return h + self.res_conv(x)
 
 class LinearAttention(nn.Module):
-    def __init__(self, dim, heads = 4, dim_head = 32):
+    def __init__(self, dim, heads = 8, dim_head = 32):
         super().__init__()
         self.heads = heads
         hidden_dim = dim_head * heads
-        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False)
+        self.to_qkv = nn.Conv2d(dim, hidden_dim, 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, qkv=3)
+        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)
         k = k.softmax(dim=-1)
         context = torch.einsum('bhdn,bhen->bhde', k, v)
@@ -358,7 +349,7 @@ class GaussianDiffusion(nn.Module):
 
     @torch.no_grad()
     def sample(self, image_size, batch_size = 16):
-        return self.p_sample_loop((batch_size, 3, image_size, image_size))
+        return self.p_sample_loop((16, 3, image_size, image_size))
 
     @torch.no_grad()
     def interpolate(self, x1, x2, t = None, lam = 0.5):
@@ -443,16 +434,13 @@ class Trainer(object):
         train_lr = 2e-5,
         train_num_steps = 100000,
         gradient_accumulate_every = 2,
-        fp16 = False,
+        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
@@ -475,7 +463,7 @@ class Trainer(object):
         self.ema_model.load_state_dict(self.model.state_dict())
 
     def step_ema(self):
-        if self.step < self.step_start_ema:
+        if self.step < 2000:
             self.reset_parameters()
             return
         self.ema.update_model_average(self.ema_model, self.model)
@@ -511,12 +499,10 @@ class Trainer(object):
             if self.step % UPDATE_EMA_EVERY == 0:
                 self.step_ema()
 
-            if self.step != 0 and self.step % SAVE_AND_SAMPLE_EVERY == 0:
+            if self.step % SAVE_AND_SAMPLE_EVERY == 0:
                 milestone = self.step // SAVE_AND_SAMPLE_EVERY
-                batches = num_to_groups(36, self.batch_size)
+                all_images = self.ema_model.p_sample_loop((64, 3, self.image_size, self.image_size))
-                all_images_list = list(map(lambda n: self.ema_model.sample(self.image_size, batch_size=n), batches))
+                utils.save_image(all_images, f'./sample-{milestone}.png', nrow=8)
-                all_images = torch.cat(all_images_list, dim=0)
-                utils.save_image(all_images, 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.2.3',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',
@@ -17,7 +17,7 @@ setup(
     'einops',
     'numpy',
     'pillow',
-    'torch',
+    'torch>=1.6',
     'torchvision',
     'tqdm'
   ],