[Fix] The interface multiscale_output is defined but not used

mmseg/models/backbones/hrnet.py

@@ -218,26 +218,41 @@ def forward(self, x):
 class HRNet(BaseModule):
     """HRNet backbone.
 
-    High-Resolution Representations for Labeling Pixels and Regions
-    arXiv: https://arxiv.org/abs/1904.04514
+    `High-Resolution Representations for Labeling Pixels and Regions
+    arXiv: <https://arxiv.org/abs/1904.04514>`_.
 
     Args:
-        extra (dict): detailed configuration for each stage of HRNet.
+        extra (dict): Detailed configuration for each stage of HRNet.
+            There must be 4 stages, the configuration for each stage must have
+            5 keys:
+
+                - num_modules (int): The number of HRModule in this stage.
+                - num_branches (int): The number of branches in the HRModule.
+                - block (str): The type of convolution block.
+                - num_blocks (tuple): The number of blocks in each branch.
+                    The length must be equal to num_branches.
+                - num_channels (tuple): The number of channels in each branch.
+                    The length must be equal to num_branches.
         in_channels (int): Number of input image channels. Normally 3.
-        conv_cfg (dict): dictionary to construct and config conv layer.
-        norm_cfg (dict): dictionary to construct and config norm layer.
+        conv_cfg (dict): Dictionary to construct and config conv layer.
+            Default: None.
+        norm_cfg (dict): Dictionary to construct and config norm layer.
+            Use `BN` by default.
         norm_eval (bool): Whether to set norm layers to eval mode, namely,
             freeze running stats (mean and var). Note: Effect on Batch Norm
-            and its variants only.
+            and its variants only. Default: False.
         with_cp (bool): Use checkpoint or not. Using checkpoint will save some
-            memory while slowing down the training speed.
+            memory while slowing down the training speed. Default: False.
         frozen_stages (int): Stages to be frozen (stop grad and set eval mode).
             -1 means not freezing any parameters. Default: -1.
-        zero_init_residual (bool): whether to use zero init for last norm layer
-            in resblocks to let them behave as identity.
-        pretrained (str, optional): model pretrained path. Default: None
+        zero_init_residual (bool): Whether to use zero init for last norm layer
+            in resblocks to let them behave as identity. Default: False.
+        multiscale_output (bool): Whether to output multi-level features
+            produced by multiple branches. If False, only the first level
+            feature will be output. Default: True.
+        pretrained (str, optional): Model pretrained path. Default: None.
         init_cfg (dict or list[dict], optional): Initialization config dict.
-            Default: None
+            Default: None.
 
     Example:
         >>> from mmseg.models import HRNet
@@ -290,6 +305,7 @@ def __init__(self,
                  with_cp=False,
                  frozen_stages=-1,
                  zero_init_residual=False,
+                 multiscale_output=True,
                  pretrained=None,
                  init_cfg=None):
         super(HRNet, self).__init__(init_cfg)
@@ -299,7 +315,7 @@ def __init__(self,
         assert not (init_cfg and pretrained), \
             'init_cfg and pretrained cannot be setting at the same time'
         if isinstance(pretrained, str):
-            warnings.warn('DeprecationWarning: pretrained is a deprecated, '
+            warnings.warn('DeprecationWarning: pretrained is deprecated, '
                           'please use "init_cfg" instead')
             self.init_cfg = dict(type='Pretrained', checkpoint=pretrained)
         elif pretrained is None:
@@ -314,6 +330,16 @@ def __init__(self,
         else:
             raise TypeError('pretrained must be a str or None')
 
+        # Assert configurations of 4 stages are in extra
+        assert 'stage1' in extra and 'stage2' in extra \
+               and 'stage3' in extra and 'stage4' in extra
+        # Assert whether the length of `num_blocks` and `num_channels` are
+        # equal to `num_branches`
+        for i in range(4):
+            cfg = extra[f'stage{i + 1}']
+            assert len(cfg['num_blocks']) == cfg['num_branches'] and \
+                   len(cfg['num_channels']) == cfg['num_branches']
+
         self.extra = extra
         self.conv_cfg = conv_cfg
         self.norm_cfg = norm_cfg
@@ -391,7 +417,7 @@ def __init__(self,
         self.transition3 = self._make_transition_layer(pre_stage_channels,
                                                        num_channels)
         self.stage4, pre_stage_channels = self._make_stage(
-            self.stage4_cfg, num_channels)
+            self.stage4_cfg, num_channels, multiscale_output=multiscale_output)
 
         self._freeze_stages()
 

tests/test_models/test_backbones/test_hrnet.py

@@ -1,12 +1,59 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import pytest
+import torch
 from mmcv.utils.parrots_wrapper import _BatchNorm
 
-from mmseg.models.backbones import HRNet
+from mmseg.models.backbones.hrnet import HRModule, HRNet
+from mmseg.models.backbones.resnet import BasicBlock, Bottleneck
 
 
-def test_hrnet_backbone():
-    # Test HRNET with two stage frozen
+@pytest.mark.parametrize('block', [BasicBlock, Bottleneck])
+def test_hrmodule(block):
+    # Test multiscale forward
+    num_channles = (32, 64)
+    in_channels = [c * block.expansion for c in num_channles]
+    hrmodule = HRModule(
+        num_branches=2,
+        blocks=block,
+        in_channels=in_channels,
+        num_blocks=(4, 4),
+        num_channels=num_channles,
+    )
+
+    feats = [
+        torch.randn(1, in_channels[0], 64, 64),
+        torch.randn(1, in_channels[1], 32, 32)
+    ]
+    feats = hrmodule(feats)
+
+    assert len(feats) == 2
+    assert feats[0].shape == torch.Size([1, in_channels[0], 64, 64])
+    assert feats[1].shape == torch.Size([1, in_channels[1], 32, 32])
+
+    # Test single scale forward
+    num_channles = (32, 64)
+    in_channels = [c * block.expansion for c in num_channles]
+    hrmodule = HRModule(
+        num_branches=2,
+        blocks=block,
+        in_channels=in_channels,
+        num_blocks=(4, 4),
+        num_channels=num_channles,
+        multiscale_output=False,
+    )
+
+    feats = [
+        torch.randn(1, in_channels[0], 64, 64),
+        torch.randn(1, in_channels[1], 32, 32)
+    ]
+    feats = hrmodule(feats)
 
+    assert len(feats) == 1
+    assert feats[0].shape == torch.Size([1, in_channels[0], 64, 64])
+
+
+def test_hrnet_backbone():
+    # only have 3 stages
     extra = dict(
         stage1=dict(
             num_modules=1,
@@ -25,13 +72,46 @@ def test_hrnet_backbone():
             num_branches=3,
             block='BASIC',
             num_blocks=(4, 4, 4),
-            num_channels=(32, 64, 128)),
-        stage4=dict(
-            num_modules=3,
-            num_branches=4,
-            block='BASIC',
-            num_blocks=(4, 4, 4, 4),
-            num_channels=(32, 64, 128, 256)))
+            num_channels=(32, 64, 128)))
+
+    with pytest.raises(AssertionError):
+        # HRNet now only support 4 stages
+        HRNet(extra=extra)
+    extra['stage4'] = dict(
+        num_modules=3,
+        num_branches=3,  # should be 4
+        block='BASIC',
+        num_blocks=(4, 4, 4, 4),
+        num_channels=(32, 64, 128, 256))
+
+    with pytest.raises(AssertionError):
+        # len(num_blocks) should equal num_branches
+        HRNet(extra=extra)
+
+    extra['stage4']['num_branches'] = 4
+
+    # Test hrnetv2p_w32
+    model = HRNet(extra=extra)
+    model.init_weights()
+    model.train()
+
+    imgs = torch.randn(1, 3, 256, 256)
+    feats = model(imgs)
+    assert len(feats) == 4
+    assert feats[0].shape == torch.Size([1, 32, 64, 64])
+    assert feats[3].shape == torch.Size([1, 256, 8, 8])
+
+    # Test single scale output
+    model = HRNet(extra=extra, multiscale_output=False)
+    model.init_weights()
+    model.train()
+
+    imgs = torch.randn(1, 3, 256, 256)
+    feats = model(imgs)
+    assert len(feats) == 1
+    assert feats[0].shape == torch.Size([1, 32, 64, 64])
+
+    # Test HRNET with two stage frozen
     frozen_stages = 2
     model = HRNet(extra, frozen_stages=frozen_stages)
     model.init_weights()