Add MMSyncBN in registry (#420)

* Add MMSyncBN in registery

* skip mmsyncbn test

mmcv/ops/sync_bn.py

@@ -6,6 +6,7 @@ from torch.autograd.function import once_differentiable
 from torch.nn.modules.module import Module
 from torch.nn.parameter import Parameter
 
+from mmcv.cnn import NORM_LAYERS
 from ..utils import ext_loader
 
 ext_module = ext_loader.load_ext('_ext', [
@@ -109,9 +110,8 @@ class SyncBatchNormFunction(Function):
             None, None, None, None
 
 
-if dist.is_available():
-
-    class SyncBatchNorm(Module):
+@NORM_LAYERS.register_module(name='MMSyncBN')
+class SyncBatchNorm(Module):
 
     def __init__(self,
                  num_features,
@@ -119,14 +119,14 @@ if dist.is_available():
                  momentum=0.1,
                  affine=True,
                  track_running_stats=True,
-                     group=dist.group.WORLD):
+                 group=None):
         super(SyncBatchNorm, self).__init__()
         self.num_features = num_features
         self.eps = eps
         self.momentum = momentum
         self.affine = affine
         self.track_running_stats = track_running_stats
-            self.group = group
+        self.group = dist.group.WORLD if group is None else group
         self.group_size = dist.get_world_size(group)
         if self.affine:
             self.weight = Parameter(torch.Tensor(num_features))
@@ -177,8 +177,8 @@ if dist.is_available():
 
         if self.training or not self.track_running_stats:
             return SyncBatchNormFunction.apply(input, self.running_mean,
-                                                   self.running_var,
-                                                   self.weight, self.bias,
+                                               self.running_var, self.weight,
+                                               self.bias,
                                                exponential_average_factor,
                                                self.eps, self.group,
                                                self.group_size)
@@ -196,12 +196,3 @@ if dist.is_available():
         s += f'track_running_stats={self.track_running_stats}, '
         s += f'group_size={self.group_size})'
         return s
-
-else:
-
-    class SyncBatchNorm(Module):
-
-        def __init__(self, *args, **kwargs):
-            raise NotImplementedError(
-                'SyncBatchNorm is not supported in this OS since the '
-                'distributed package is not available')

tests/test_cnn/test_build_layers.py

@@ -135,6 +135,8 @@ def test_build_norm_layer():
         'IN3d': 'in',
     }
     for type_name, module in NORM_LAYERS.module_dict.items():
+        if type_name == 'MMSyncBN':  # skip MMSyncBN
+            continue
         for postfix in ['_test', 1]:
             cfg = dict(type=type_name)
             if type_name == 'GN':