higher epsilon for fp16 in layernorm

README.md

@@ -1,4 +1,4 @@
-<img src="./denoising-diffusion.png" width="500px"></img>
+<img src="./images/denoising-diffusion.png" width="500px"></img>
 
 ## Denoising Diffusion Probabilistic Model, in Pytorch
 
@@ -10,7 +10,7 @@ Youtube AI Educators - <a href="https://www.youtube.com/watch?v=W-O7AZNzbzQ">Yan
 
 <a href="https://huggingface.co/blog/annotated-diffusion">Annotated code</a> by Research Scientists / Engineers from <a href="https://huggingface.co/">🤗 Huggingface</a>
 
-<img src="./sample.png" width="500px"><img>
+<img src="./images/sample.png" width="500px"><img>
 
 [![PyPI version](https://badge.fury.io/py/denoising-diffusion-pytorch.svg)](https://badge.fury.io/py/denoising-diffusion-pytorch)
 

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -89,16 +89,15 @@ def Downsample(dim, dim_out = None):
     return nn.Conv2d(dim, default(dim_out, dim), 4, 2, 1)
 
 class LayerNorm(nn.Module):
-    def __init__(self, dim, eps = 1e-5):
+    def __init__(self, dim):
         super().__init__()
-        self.eps = eps
         self.g = nn.Parameter(torch.ones(1, dim, 1, 1))
-        self.b = nn.Parameter(torch.zeros(1, dim, 1, 1))
 
     def forward(self, x):
+        eps = 1e-5 if x.dtype == torch.float32 else 1e-3
         var = torch.var(x, dim = 1, unbiased = False, keepdim = True)
         mean = torch.mean(x, dim = 1, keepdim = True)
-        return (x - mean) / (var + self.eps).sqrt() * self.g + self.b
+        return (x - mean) * (var + eps).rsqrt() * self.g
 
 class PreNorm(nn.Module):
     def __init__(self, dim, fn):
@@ -211,6 +210,8 @@ class LinearAttention(nn.Module):
         k = k.softmax(dim = -1)
 
         q = q * self.scale
+        v = v / (h * w)
+
         context = torch.einsum('b h d n, b h e n -> b h d e', k, v)
 
         out = torch.einsum('b h d e, b h d n -> b h e n', context, q)
@@ -479,7 +480,7 @@ class GaussianDiffusion(nn.Module):
 
     def predict_noise_from_start(self, x_t, t, x0):
         return (
-            (x0 - extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t) / \
+            (extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - x0) / \
             extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
         )
 
@@ -719,6 +720,7 @@ class Trainer(object):
         self.ds = Dataset(folder, self.image_size, augment_horizontal_flip = augment_horizontal_flip, convert_image_to = convert_image_to)
         dl = DataLoader(self.ds, batch_size = train_batch_size, shuffle = True, pin_memory = True, num_workers = cpu_count())
 
+        dl = self.accelerator.prepare(dl)
         self.dl = cycle(dl)
 
         # optimizer
@@ -739,7 +741,7 @@ class Trainer(object):
 
         # prepare model, dataloader, optimizer with accelerator
 
-        self.model, self.dl, self.opt = self.accelerator.prepare(self.model, self.dl, self.opt)
+        self.model, self.opt = self.accelerator.prepare(self.model, self.opt)
 
     def save(self, milestone):
         if not self.accelerator.is_local_main_process:

setup.py

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