Merge pull request #61 from hellock/weight-init

Add caffe2_xavier_init()

mmcv/cnn/__init__.py

@@ -2,10 +2,10 @@ from .alexnet import AlexNet
 from .vgg import VGG, make_vgg_layer
 from .resnet import ResNet, make_res_layer
 from .weight_init import (constant_init, xavier_init, normal_init,
-                          uniform_init, kaiming_init)
+                          uniform_init, kaiming_init, caffe2_xavier_init)
 
 __all__ = [
     'AlexNet', 'VGG', 'make_vgg_layer', 'ResNet', 'make_res_layer',
     'constant_init', 'xavier_init', 'normal_init', 'uniform_init',
-    'kaiming_init'
+    'kaiming_init', 'caffe2_xavier_init'
 ]

mmcv/cnn/weight_init.py

@@ -30,6 +30,7 @@ def uniform_init(module, a=0, b=1, bias=0):
 
 
 def kaiming_init(module,
+                 a=0,
                  mode='fan_out',
                  nonlinearity='relu',
                  bias=0,
@@ -37,9 +38,20 @@ def kaiming_init(module,
     assert distribution in ['uniform', 'normal']
     if distribution == 'uniform':
         nn.init.kaiming_uniform_(
-            module.weight, mode=mode, nonlinearity=nonlinearity)
+            module.weight, a=a, mode=mode, nonlinearity=nonlinearity)
     else:
         nn.init.kaiming_normal_(
-            module.weight, mode=mode, nonlinearity=nonlinearity)
+            module.weight, a=a, mode=mode, nonlinearity=nonlinearity)
     if hasattr(module, 'bias') and module.bias is not None:
         nn.init.constant_(module.bias, bias)
+
+
+def caffe2_xavier_init(module, bias=0):
+    # `XavierFill` in Caffe2 corresponds to `kaiming_uniform_` in PyTorch
+    # Acknowledgment to FAIR's internal code
+    kaiming_init(
+        module,
+        a=1,
+        mode='fan_in',
+        nonlinearity='leaky_relu',
+        distribution='uniform')