add weight standardization prior to groupnorm, lessen cosine sim attention scale to 10 for fp16

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -97,11 +97,16 @@ class WeightStandardizedConv2d(nn.Conv2d):
         eps = 1e-5 if x.dtype == torch.float32 else 1e-3
 
         weight = self.weight
-        mean = reduce(weight, 'o ... -> o 1 1 1', 'mean')
-        var = reduce(weight, 'o ... -> o 1 1 1', partial(torch.var, unbiased = False))
-        normalized_weight = (weight - mean) * (var + eps).rsqrt()
+        flattened_weights = rearrange(weight, 'o ... -> o (...)')
 
-        return F.conv2d(x, normalized_weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
+        mean = reduce(weight, 'o ... -> o 1 1 1', 'mean')
+
+        var = torch.var(flattened_weights, dim = -1, unbiased = False)
+        var = rearrange(var, 'o -> o 1 1 1')
+
+        weight = (weight - mean) * (var + eps).rsqrt()
+
+        return F.conv2d(x, weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
 
 class LayerNorm(nn.Module):
     def __init__(self, dim):

setup.py

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