Fix group norm layer

Author: roytseng-tw

lib/nn/functional.py

@@ -2,13 +2,17 @@
 
 
 def group_norm(x, num_groups, weight=None, bias=None, eps=1e-5):
-    N, C, H, W = x.size()
-    assert C % num_groups == 0, "input channel dimension must divisible by number of groups"
-    x = x.view(N, num_groups, -1)
+    input_shape = x.shape
+    ndim = len(input_shape)
+    N, C = input_shape[:2]
+    G = num_groups
+    assert C % G == 0, "input channel dimension must divisible by number of groups"
+    x = x.view(N, G, -1)
     mean = x.mean(-1, keepdim=True)
     var = x.var(-1, keepdim=True)
     x = (x - mean) / (var + eps).sqrt()
-    x = x.view(N, C, H, W)
-    if weight is not None:  # affine=True
-        return x * weight + bias
-    return x
+    x = x.view(input_shape)
+    view_shape = (1, -1) + (1,) * (ndim - 2)
+    if weight is not None:
+        return x * weight.view(view_shape) + bias.view(view_shape)
+    return x

lib/nn/modules/normalization.py

@@ -8,7 +8,7 @@
 
 class GroupNorm(nn.Module):
     def __init__(self, num_groups, num_channels, eps=1e-5, affine=True):
-        super(GroupNorm, self).__init__()
+        super().__init__()
         self.num_groups = num_groups
         self.num_channels = num_channels
         self.eps = eps
@@ -28,7 +28,7 @@ def reset_parameters(self):
 
     def forward(self, x):
         return myF.group_norm(
-            x, self.num_groups, self.weight, self.bias, self.eps)
+            x, self.num_groups, self.weight, self.bias, self.eps
         )
 
     def extra_repr(self):