SparseGroupNorm

sparseconvnet/__init__.py

@@ -8,7 +8,7 @@ forward_pass_multiplyAdd_count = 0
 forward_pass_hidden_states = 0
 from .activations import Tanh, Sigmoid, ReLU, LeakyReLU, ELU, SELU, BatchNormELU
 from .averagePooling import AveragePooling
-from .batchNormalization import BatchNormalization, BatchNormReLU, BatchNormLeakyReLU, MeanOnlyBNLeakyReLU
+from .batchNormalization import BatchNormalization, BatchNormReLU, BatchNormLeakyReLU, MeanOnlyBNLeakyReLU, SparseGroupNorm
 from .classificationTrainValidate import ClassificationTrainValidate
 from .convolution import Convolution
 from .deconvolution import Deconvolution

sparseconvnet/batchNormalization.py

@@ -208,3 +208,11 @@ class MeanOnlyBNLeakyReLU(Module):
     def __repr__(self):
         s = 'MeanOnlyBatchNorm(' + str(self.nPlanes) + ',momentum=' + str(self.momentum) + ',leakiness=' + str(self.leakiness) + ')'
         return s
+
+
+class SparseGroupNorm(torch.nn.GroupNorm):
+    def forward(self,x):
+        return scn.SparseConvNetTensor(
+            super().forward(x.features),
+            x.metadata,
+            x.spatial_size)

sparseconvnet/utils.py

@@ -335,3 +335,10 @@ class VerboseIdentity(torch.nn.Module):
     def forward(self, x):
         print(x)
         return x
+
+class SparseGroupNorm(torch.nn.GroupNorm):
+    def forward(self,x):
+        return scn.SparseConvNetTensor(
+            super().forward(x.features),
+            x.metadata,
+            x.spatial_size)