MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / OriginGroupNormFunc

Class OriginGroupNormFunc

imperative/python/test/unit/module/test_module.py:730–756  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

728 bias_np = np.random.uniform(-0.5, 0.5, (num_channels))
729
730 class OriginGroupNormFunc(Module):
731 def __init__(self, eps=1e-5, affine=True, **kwargs):
732 super().__init__(**kwargs)
733 assert num_channels % num_groups == 0
734 self.num_groups = num_groups
735 self.num_channels = num_channels
736 self.eps = eps
737 self.affine = affine
738 if self.affine:
739 self.weight = Parameter(weight_np)
740 self.bias = Parameter(bias_np)
741 else:
742 self.weight = None
743 self.bias = None
744
745 def forward(self, x):
746 N, C, H, W = x.shape
747 x = x.reshape(N, self.num_groups, -1)
748 mean = x.mean(axis=2, keepdims=True)
749 var = (x * x).mean(axis=2, keepdims=True) - mean * mean
750 x = (x - mean) / F.sqrt(var + self.eps)
751 x = x.reshape(N, C, H, W)
752 if self.affine:
753 x = self.weight.reshape(1, -1, 1, 1) * x + self.bias.reshape(
754 1, -1, 1, 1
755 )
756 return x
757
758 inp = np.random.uniform(-0.5, 0.5, (2, num_channels, 10, 16)).astype("float32")
759 mge_inp = Tensor(inp)

Callers 1

test_grou_normFunction · 0.85

Calls

no outgoing calls

Tested by 1

test_grou_normFunction · 0.68