Support to pass args to cnn bricks (#277)

mmcv/cnn/bricks/upsample.py

@@ -47,7 +47,7 @@ class PixelShufflePack(nn.Module):
         return x
 
 
-def build_upsample_layer(cfg):
+def build_upsample_layer(cfg, *args, **kwargs):
     """Build upsample layer.
 
     Args:
@@ -56,6 +56,10 @@ def build_upsample_layer(cfg):
             - scale_factor (int): Upsample ratio, which is not applicable to
                 deconv.
             - layer args: Args needed to instantiate a upsample layer.
+        args (argument list): Arguments passed to the `__init__`
+            method of the corresponding conv layer.
+        kwargs (keyword arguments): Keyword arguments passed to the `__init__`
+            method of the corresponding conv layer.
 
     Returns:
         nn.Module: Created upsample layer.
@@ -75,5 +79,5 @@ def build_upsample_layer(cfg):
 
     if upsample is nn.Upsample:
         cfg_['mode'] = layer_type
-    layer = upsample(**cfg_)
+    layer = upsample(*args, **kwargs, **cfg_)
     return layer

tests/test_cnn/test_build_layers.py

@@ -226,6 +226,22 @@ def test_upsample_layer():
     layer = build_upsample_layer(cfg)
     assert isinstance(layer, nn.ConvTranspose2d)
 
+    cfg = dict(type='deconv')
+    kwargs = dict(in_channels=3, out_channels=3, kernel_size=3, stride=2)
+    layer = build_upsample_layer(cfg, **kwargs)
+    assert isinstance(layer, nn.ConvTranspose2d)
+    assert layer.in_channels == kwargs['in_channels']
+    assert layer.out_channels == kwargs['out_channels']
+    assert layer.kernel_size == (kwargs['kernel_size'], kwargs['kernel_size'])
+    assert layer.stride == (kwargs['stride'], kwargs['stride'])
+
+    layer = build_upsample_layer(cfg, 3, 3, 3, 2)
+    assert isinstance(layer, nn.ConvTranspose2d)
+    assert layer.in_channels == kwargs['in_channels']
+    assert layer.out_channels == kwargs['out_channels']
+    assert layer.kernel_size == (kwargs['kernel_size'], kwargs['kernel_size'])
+    assert layer.stride == (kwargs['stride'], kwargs['stride'])
+
     cfg = dict(
         type='pixel_shuffle',
         in_channels=3,