Add HRNet feature extraction, fix senet type, lower feature testing res to 96x96

rwightman · rwightman · commit c9d54bc1c31f · 2020-07-21T17:39:29.000-07:00
diff --git a/tests/test_models.py b/tests/test_models.py
@@ -109,12 +109,13 @@ def test_model_forward_torchscript(model_name, batch_size):
 
 
 EXCLUDE_FEAT_FILTERS = [
-    'hrnet*',  '*pruned*',  # hopefully fix at some point
+    '*pruned*',  # hopefully fix at some point
 ]
 if 'GITHUB_ACTIONS' in os.environ and 'Linux' in platform.system():
     # GitHub Linux runner is slower and hits memory limits sooner than MacOS, exclude bigger models
     EXCLUDE_FEAT_FILTERS += ['*resnext101_32x32d']
 
+
 @pytest.mark.timeout(120)
 @pytest.mark.parametrize('model_name', list_models(exclude_filters=EXCLUDE_FILTERS + EXCLUDE_FEAT_FILTERS))
 @pytest.mark.parametrize('batch_size', [1])
@@ -124,7 +125,7 @@ def test_model_forward_features(model_name, batch_size):
     model.eval()
     expected_channels = model.feature_info.channels()
     assert len(expected_channels) >= 4  # all models here should have at least 4 feature levels by default, some 5 or 6
-    input_size = (3, 128, 128)  # jit compile is already a bit slow and we've tested normal res already...
+    input_size = (3, 96, 96)  # jit compile is already a bit slow and we've tested normal res already...
     outputs = model(torch.randn((batch_size, *input_size)))
     assert len(expected_channels) == len(outputs)
     for e, o in zip(expected_channels, outputs):
diff --git a/timm/models/hrnet.py b/timm/models/hrnet.py
@@ -8,17 +8,15 @@
   Written by Bin Xiao (Bin.Xiao@microsoft.com)
   Modified by Ke Sun (sunk@mail.ustc.edu.cn)
 """
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import logging
+from typing import List
 
+import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
 from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+from .features import FeatureInfo
 from .helpers import build_model_with_cfg
 from .layers import SelectAdaptivePool2d
 from .registry import register_model
@@ -403,32 +401,23 @@ def __init__(self, num_branches, blocks, num_blocks, num_inchannels,
         self.branches = self._make_branches(
             num_branches, blocks, num_blocks, num_channels)
         self.fuse_layers = self._make_fuse_layers()
-        self.relu = nn.ReLU(False)
+        self.fuse_act = nn.ReLU(False)
 
     def _check_branches(self, num_branches, blocks, num_blocks, num_inchannels, num_channels):
+        error_msg = ''
         if num_branches != len(num_blocks):
-            error_msg = 'NUM_BRANCHES({}) <> NUM_BLOCKS({})'.format(
-                num_branches, len(num_blocks))
-            logger.error(error_msg)
-            raise ValueError(error_msg)
-
-        if num_branches != len(num_channels):
-            error_msg = 'NUM_BRANCHES({}) <> NUM_CHANNELS({})'.format(
-                num_branches, len(num_channels))
+            error_msg = 'NUM_BRANCHES({}) <> NUM_BLOCKS({})'.format(num_branches, len(num_blocks))
+        elif num_branches != len(num_channels):
+            error_msg = 'NUM_BRANCHES({}) <> NUM_CHANNELS({})'.format(num_branches, len(num_channels))
+        elif num_branches != len(num_inchannels):
+            error_msg = 'NUM_BRANCHES({}) <> NUM_INCHANNELS({})'.format(num_branches, len(num_inchannels))
+        if error_msg:
             logger.error(error_msg)
             raise ValueError(error_msg)
 
-        if num_branches != len(num_inchannels):
-            error_msg = 'NUM_BRANCHES({}) <> NUM_INCHANNELS({})'.format(
-                num_branches, len(num_inchannels))
-            logger.error(error_msg)
-            raise ValueError(error_msg)
-
-    def _make_one_branch(self, branch_index, block, num_blocks, num_channels,
-                         stride=1):
+    def _make_one_branch(self, branch_index, block, num_blocks, num_channels, stride=1):
         downsample = None
-        if stride != 1 or \
-                self.num_inchannels[branch_index] != num_channels[branch_index] * block.expansion:
+        if stride != 1 or self.num_inchannels[branch_index] != num_channels[branch_index] * block.expansion:
             downsample = nn.Sequential(
                 nn.Conv2d(
                     self.num_inchannels[branch_index], num_channels[branch_index] * block.expansion,
@@ -489,22 +478,22 @@ def _make_fuse_layers(self):
     def get_num_inchannels(self):
         return self.num_inchannels
 
-    def forward(self, x):
+    def forward(self, x: List[torch.Tensor]):
         if self.num_branches == 1:
             return [self.branches[0](x[0])]
 
-        for i in range(self.num_branches):
-            x[i] = self.branches[i](x[i])
+        for i, branch in enumerate(self.branches):
+            x[i] = branch(x[i])
 
         x_fuse = []
-        for i in range(len(self.fuse_layers)):
-            y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
+        for i, fuse_outer in enumerate(self.fuse_layers):
+            y = x[0] if i == 0 else fuse_outer[0](x[0])
             for j in range(1, self.num_branches):
                 if i == j:
                     y = y + x[j]
                 else:
-                    y = y + self.fuse_layers[i][j](x[j])
-            x_fuse.append(self.relu(y))
+                    y = y + fuse_outer[j](x[j])
+            x_fuse.append(self.fuse_act(y))
 
         return x_fuse
 
@@ -517,17 +506,18 @@ def forward(self, x):
 
 class HighResolutionNet(nn.Module):
 
-    def __init__(self, cfg, in_chans=3, num_classes=1000, global_pool='avg', drop_rate=0.0):
+    def __init__(self, cfg, in_chans=3, num_classes=1000, global_pool='avg', drop_rate=0.0, head='classification'):
         super(HighResolutionNet, self).__init__()
         self.num_classes = num_classes
         self.drop_rate = drop_rate
 
         stem_width = cfg['STEM_WIDTH']
         self.conv1 = nn.Conv2d(in_chans, stem_width, kernel_size=3, stride=2, padding=1, bias=False)
         self.bn1 = nn.BatchNorm2d(stem_width, momentum=_BN_MOMENTUM)
+        self.act1 = nn.ReLU(inplace=True)
         self.conv2 = nn.Conv2d(stem_width, 64, kernel_size=3, stride=2, padding=1, bias=False)
         self.bn2 = nn.BatchNorm2d(64, momentum=_BN_MOMENTUM)
-        self.relu = nn.ReLU(inplace=True)
+        self.act2 = nn.ReLU(inplace=True)
 
         self.stage1_cfg = cfg['STAGE1']
         num_channels = self.stage1_cfg['NUM_CHANNELS'][0]
@@ -557,31 +547,49 @@ def __init__(self, cfg, in_chans=3, num_classes=1000, global_pool='avg', drop_ra
         self.transition3 = self._make_transition_layer(pre_stage_channels, num_channels)
         self.stage4, pre_stage_channels = self._make_stage(self.stage4_cfg, num_channels, multi_scale_output=True)
 
-        # Classification Head
-        self.num_features = 2048
-        self.incre_modules, self.downsamp_modules, self.final_layer = self._make_head(pre_stage_channels)
-        self.global_pool = SelectAdaptivePool2d(pool_type=global_pool)
-        self.classifier = nn.Linear(self.num_features * self.global_pool.feat_mult(), num_classes)
+        self.head = head
+        self.head_channels = None  # set if _make_head called
+        if head == 'classification':
+            # Classification Head
+            self.num_features = 2048
+            self.incre_modules, self.downsamp_modules, self.final_layer = self._make_head(pre_stage_channels)
+            self.global_pool = SelectAdaptivePool2d(pool_type=global_pool)
+            self.classifier = nn.Linear(self.num_features * self.global_pool.feat_mult(), num_classes)
+        elif head == 'incre':
+            self.num_features = 2048
+            self.incre_modules, _, _ = self._make_head(pre_stage_channels, True)
+        else:
+            self.incre_modules = None
+            self.num_features = 256
+
+        curr_stride = 2
+        # module names aren't actually valid here, hook or FeatureNet based extraction would not work
+        self.feature_info = [dict(num_chs=64, reduction=curr_stride, module='stem')]
+        for i, c in enumerate(self.head_channels if self.head_channels else num_channels):
+            curr_stride *= 2
+            c = c * 4 if self.head_channels else c  # head block expansion factor of 4
+            self.feature_info += [dict(num_chs=c, reduction=curr_stride, module=f'stage{i + 1}')]
 
         self.init_weights()
 
-    def _make_head(self, pre_stage_channels):
+    def _make_head(self, pre_stage_channels, incre_only=False):
         head_block = Bottleneck
-        head_channels = [32, 64, 128, 256]
+        self.head_channels = [32, 64, 128, 256]
 
         # Increasing the #channels on each resolution
         # from C, 2C, 4C, 8C to 128, 256, 512, 1024
         incre_modules = []
         for i, channels in enumerate(pre_stage_channels):
-            incre_modules.append(
-                self._make_layer(head_block, channels, head_channels[i], 1, stride=1))
+            incre_modules.append(self._make_layer(head_block, channels, self.head_channels[i], 1, stride=1))
         incre_modules = nn.ModuleList(incre_modules)
+        if incre_only:
+            return incre_modules, None, None
 
         # downsampling modules
         downsamp_modules = []
         for i in range(len(pre_stage_channels) - 1):
-            in_channels = head_channels[i] * head_block.expansion
-            out_channels = head_channels[i + 1] * head_block.expansion
+            in_channels = self.head_channels[i] * head_block.expansion
+            out_channels = self.head_channels[i + 1] * head_block.expansion
             downsamp_module = nn.Sequential(
                 nn.Conv2d(
                     in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=2, padding=1),
@@ -593,7 +601,7 @@ def _make_head(self, pre_stage_channels):
 
         final_layer = nn.Sequential(
             nn.Conv2d(
-                in_channels=head_channels[3] * head_block.expansion,
+                in_channels=self.head_channels[3] * head_block.expansion,
                 out_channels=self.num_features, kernel_size=1, stride=1, padding=0
             ),
             nn.BatchNorm2d(self.num_features, momentum=_BN_MOMENTUM),
@@ -655,11 +663,7 @@ def _make_stage(self, layer_config, num_inchannels, multi_scale_output=True):
         modules = []
         for i in range(num_modules):
             # multi_scale_output is only used last module
-            if not multi_scale_output and i == num_modules - 1:
-                reset_multi_scale_output = False
-            else:
-                reset_multi_scale_output = True
-
+            reset_multi_scale_output = multi_scale_output or i < num_modules - 1
             modules.append(HighResolutionModule(
                 num_branches, block, num_blocks, num_inchannels, num_channels, fuse_method, reset_multi_scale_output)
             )
@@ -688,40 +692,35 @@ def reset_classifier(self, num_classes, global_pool='avg'):
         else:
             self.classifier = nn.Identity()
 
+    def stages(self, x) -> List[torch.Tensor]:
+        x = self.layer1(x)
+
+        xl = [t(x) for i, t in enumerate(self.transition1)]
+        yl = self.stage2(xl)
+
+        xl = [t(yl[-1]) if not isinstance(t, nn.Identity) else yl[i] for i, t in enumerate(self.transition2)]
+        yl = self.stage3(xl)
+
+        xl = [t(yl[-1]) if not isinstance(t, nn.Identity) else yl[i] for i, t in enumerate(self.transition3)]
+        yl = self.stage4(xl)
+        return yl
+
     def forward_features(self, x):
+        # Stem
         x = self.conv1(x)
         x = self.bn1(x)
-        x = self.relu(x)
+        x = self.act1(x)
         x = self.conv2(x)
         x = self.bn2(x)
-        x = self.relu(x)
-        x = self.layer1(x)
-
-        x_list = []
-        for i in range(len(self.transition1)):
-            x_list.append(self.transition1[i](x))
-        y_list = self.stage2(x_list)
-
-        x_list = []
-        for i in range(len(self.transition2)):
-            if not isinstance(self.transition2[i], nn.Identity):
-                x_list.append(self.transition2[i](y_list[-1]))
-            else:
-                x_list.append(y_list[i])
-        y_list = self.stage3(x_list)
+        x = self.act2(x)
 
-        x_list = []
-        for i in range(len(self.transition3)):
-            if not isinstance(self.transition3[i], nn.Identity):
-                x_list.append(self.transition3[i](y_list[-1]))
-            else:
-                x_list.append(y_list[i])
-        y_list = self.stage4(x_list)
+        # Stages
+        yl = self.stages(x)
 
         # Classification Head
-        y = self.incre_modules[0](y_list[0])
-        for i in range(len(self.downsamp_modules)):
-            y = self.incre_modules[i + 1](y_list[i + 1]) + self.downsamp_modules[i](y)
+        y = self.incre_modules[0](yl[0])
+        for i, down in enumerate(self.downsamp_modules):
+            y = self.incre_modules[i + 1](yl[i + 1]) + down(y)
         y = self.final_layer(y)
         return y
 
@@ -734,10 +733,55 @@ def forward(self, x):
         return x
 
 
+class HighResolutionNetFeatures(HighResolutionNet):
+    """HighResolutionNet feature extraction
+
+    The design of HRNet makes it easy to grab feature maps, this class provides a simple wrapper to do so.
+    It would be more complicated to use the FeatureNet helpers.
+
+    The `feature_location=incre` allows grabbing increased channel count features using part of the
+    classification head. If `feature_location=''` the default HRNet features are returned. First stem
+    conv is used for stride 2 features.
+    """
+
+    def __init__(self, cfg, in_chans=3, num_classes=1000, global_pool='avg', drop_rate=0.0,
+                 feature_location='incre', out_indices=(0, 1, 2, 3, 4)):
+        assert feature_location in ('incre', '')
+        super(HighResolutionNetFeatures, self).__init__(
+            cfg, in_chans=in_chans, num_classes=num_classes, global_pool=global_pool,
+            drop_rate=drop_rate, head=feature_location)
+        self.feature_info = FeatureInfo(self.feature_info, out_indices)
+        self._out_idx = {i for i in out_indices}
+
+    def forward_features(self, x):
+        assert False, 'Not supported'
+
+    def forward(self, x) -> List[torch.tensor]:
+        out = []
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.act1(x)
+        if 0 in self._out_idx:
+            out.append(x)
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.act2(x)
+        x = self.stages(x)
+        if self.incre_modules is not None:
+            x = [incre(f) for f, incre in zip(x, self.incre_modules)]
+        for i, f in enumerate(x):
+            if i + 1 in self._out_idx:
+                out.append(f)
+        return out
+
+
 def _create_hrnet(variant, pretrained, **model_kwargs):
-    assert not model_kwargs.pop('features_only', False)  # feature extraction not figured out yet
+    model_cls = HighResolutionNet
+    if model_kwargs.pop('features_only', False):
+        model_cls = HighResolutionNetFeatures
+
     return build_model_with_cfg(
-        HighResolutionNet, variant, pretrained, default_cfg=default_cfgs[variant],
+        model_cls, variant, pretrained, default_cfg=default_cfgs[variant],
         model_cfg=cfg_cls[variant], **model_kwargs)
 
 
diff --git a/timm/models/senet.py b/timm/models/senet.py
@@ -423,14 +423,14 @@ def legacy_seresnet34(pretrained=False, **kwargs):
 
 
 @register_model
-def legacy_seresnet50(pretrained=False, num_classes=1000, in_chans=3, **kwargs):
+def legacy_seresnet50(pretrained=False, **kwargs):
     model_args = dict(
         block=SEResNetBottleneck, layers=[3, 4, 6, 3], groups=1, reduction=16, **kwargs)
     return _create_senet('seresnet50', pretrained, **model_args)
 
 
 @register_model
-def legacy_seresnet101(pretrained=False, num_classes=1000, in_chans=3, **kwargs):
+def legacy_seresnet101(pretrained=False, **kwargs):
     model_args = dict(
         block=SEResNetBottleneck, layers=[3, 4, 23, 3], groups=1, reduction=16, **kwargs)
     return _create_senet('seresnet101', pretrained, **model_args)
