fix layernorm

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -111,11 +111,23 @@ class Downsample(nn.Module):
     def forward(self, x):
         return self.conv(x)
 
+class LayerNorm(nn.Module):
+    def __init__(self, dim, eps = 1e-5):
+        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):
+        std = torch.var(x, dim = 1, unbiased = False, keepdim = True).sqrt()
+        mean = torch.mean(x, dim = 1, keepdim = True)
+        return (x - mean) / (std + self.eps) * self.g + self.b
+
 class PreNorm(nn.Module):
     def __init__(self, dim, fn):
         super().__init__()
         self.fn = fn
-        self.norm = nn.InstanceNorm2d(dim, affine = True)
+        self.norm = LayerNorm(dim)
 
     def forward(self, x):
         x = self.norm(x)
@@ -150,7 +162,6 @@ class ResnetBlock(nn.Module):
         h = self.block1(x)
 
         if exists(self.mlp):
-            print('hmmm')
             h += self.mlp(time_emb)[:, :, None, None]
 
         h = self.block2(h)

setup.py

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