[Enhance] Enhance feature extraction function.

configs/_base_/models/mobilenet_v3_large_imagenet.py

@@ -11,4 +11,6 @@
         dropout_rate=0.2,
         act_cfg=dict(type='HSwish'),
         loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
+        init_cfg=dict(
+            type='Normal', layer='Linear', mean=0., std=0.01, bias=0.),
         topk=(1, 5)))

configs/_base_/models/mobilenet_v3_small_imagenet.py

@@ -11,4 +11,6 @@
         dropout_rate=0.2,
         act_cfg=dict(type='HSwish'),
         loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
+        init_cfg=dict(
+            type='Normal', layer='Linear', mean=0., std=0.01, bias=0.),
         topk=(1, 5)))

configs/mobilenet_v3/mobilenet-v3-small_8xb32_in1k.py

@@ -17,7 +17,7 @@
 # - modify: RandomErasing use RE-M instead of RE-0
 
 _base_ = [
-    '../_base_/models/mobilenet-v3-small_8xb32_in1k.py',
+    '../_base_/models/mobilenet_v3_small_imagenet.py',
     '../_base_/datasets/imagenet_bs32_pil_resize.py',
     '../_base_/default_runtime.py'
 ]

mmcls/models/classifiers/base.py

@@ -35,13 +35,14 @@ def with_head(self):
         return hasattr(self, 'head') and self.head is not None
 
     @abstractmethod
-    def extract_feat(self, imgs):
+    def extract_feat(self, imgs, stage=None):
         pass
 
-    def extract_feats(self, imgs):
+    def extract_feats(self, imgs, stage=None):
         assert isinstance(imgs, list)
+        kwargs = {} if stage is None else {'stage': stage}
         for img in imgs:
-            yield self.extract_feat(img)
+            yield self.extract_feat(img, **kwargs)
 
     @abstractmethod
     def forward_train(self, imgs, **kwargs):

mmcls/models/classifiers/image.py

@@ -70,8 +70,12 @@ def __init__(self,
                     cfg['prob'] = cutmix_prob
                     self.augments = Augments(cfg)
 
-    def extract_feat(self, img):
+    def extract_feat(self, img, stage='neck'):
         """Directly extract features from the backbone + neck."""
+        assert stage in ['backbone', 'neck', 'pre_logits'], \
+            (f'Invalid output stage "{stage}", please choose from "backbone", '
+             '"neck" and "pre_logits"')
+
         x = self.backbone(img)
         if self.return_tuple:
             if not isinstance(x, tuple):
@@ -83,8 +87,16 @@ def extract_feat(self, img):
         else:
             if isinstance(x, tuple):
                 x = x[-1]
+        if stage == 'backbone':
+            return x
+
         if self.with_neck:
             x = self.neck(x)
+        if stage == 'neck':
+            return x
+
+        if self.with_head and hasattr(self.head, 'pre_logits'):
+            x = self.head.pre_logits(x)
         return x
 
     def forward_train(self, img, gt_label, **kwargs):
@@ -122,12 +134,12 @@ def forward_train(self, img, gt_label, **kwargs):
 
         return losses
 
-    def simple_test(self, img, img_metas):
+    def simple_test(self, img, img_metas=None, **kwargs):
         """Test without augmentation."""
         x = self.extract_feat(img)
 
         try:
-            res = self.head.simple_test(x)
+            res = self.head.simple_test(x, **kwargs)
         except TypeError as e:
             if 'not tuple' in str(e) and self.return_tuple:
                 return TypeError(

mmcls/models/heads/cls_head.py

@@ -1,4 +1,6 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import warnings
+
 import torch
 import torch.nn.functional as F
 
@@ -62,14 +64,33 @@ def forward_train(self, cls_score, gt_label, **kwargs):
         losses = self.loss(cls_score, gt_label, **kwargs)
         return losses
 
-    def simple_test(self, cls_score):
+    def pre_logits(self, x):
+        if isinstance(x, tuple):
+            x = x[-1]
+
+        warnings.warn(
+            'The input of ClsHead should be already logits. '
+            'Please modify the backbone if you want to get pre-logits feature.'
+        )
+        return x
+
+    def simple_test(self, cls_score, softmax=True, post_process=True):
         """Test without augmentation."""
         if isinstance(cls_score, tuple):
             cls_score = cls_score[-1]
         if isinstance(cls_score, list):
             cls_score = sum(cls_score) / float(len(cls_score))
-        pred = F.softmax(cls_score, dim=1) if cls_score is not None else None
-        return self.post_process(pred)
+
+        if softmax:
+            pred = (
+                F.softmax(cls_score, dim=1) if cls_score is not None else None)
+        else:
+            pred = cls_score
+
+        if post_process:
+            return self.post_process(pred)
+        else:
+            return pred
 
     def post_process(self, pred):
         on_trace = is_tracing()

mmcls/models/heads/conformer_head.py

@@ -16,7 +16,7 @@ class ConformerHead(ClsHead):
             category.
         in_channels (int): Number of channels in the input feature map.
         init_cfg (dict | optional): The extra init config of layers.
-            Defaults to use dict(type='Normal', layer='Linear', std=0.01).
+            Defaults to use ``dict(type='Normal', layer='Linear', std=0.01)``.
     """
 
     def __init__(
@@ -55,25 +55,34 @@ def init_weights(self):
         else:
             self.apply(self._init_weights)
 
-    def simple_test(self, x):
-        """Test without augmentation."""
+    def pre_logits(self, x):
         if isinstance(x, tuple):
             x = x[-1]
-        assert isinstance(x,
-                          list)  # There are two outputs in the Conformer model
+        return x
+
+    def simple_test(self, x, softmax=True, post_process=True):
+        """Test without augmentation."""
+        x = self.pre_logits(x)
+        # There are two outputs in the Conformer model
+        assert isinstance(x, list)
 
         conv_cls_score = self.conv_cls_head(x[0])
         tran_cls_score = self.trans_cls_head(x[1])
 
-        cls_score = conv_cls_score + tran_cls_score
-
-        pred = F.softmax(cls_score, dim=1) if cls_score is not None else None
-
-        return self.post_process(pred)
+        if softmax:
+            cls_score = conv_cls_score + tran_cls_score
+            pred = (
+                F.softmax(cls_score, dim=1) if cls_score is not None else None)
+            if post_process:
+                pred = self.post_process(pred)
+        else:
+            pred = [conv_cls_score, tran_cls_score]
+            if post_process:
+                pred = list(map(self.post_process, pred))
+        return pred
 
     def forward_train(self, x, gt_label):
-        if isinstance(x, tuple):
-            x = x[-1]
+        x = self.pre_logits(x)
         assert isinstance(x, list) and len(x) == 2, \
             'There should be two outputs in the Conformer model'
 

mmcls/models/heads/linear_head.py

@@ -35,20 +35,31 @@ def __init__(self,
 
         self.fc = nn.Linear(self.in_channels, self.num_classes)
 
-    def simple_test(self, x):
-        """Test without augmentation."""
+    def pre_logits(self, x):
         if isinstance(x, tuple):
             x = x[-1]
+        return x
+
+    def simple_test(self, x, softmax=True, post_process=True):
+        """Test without augmentation."""
+        x = self.pre_logits(x)
         cls_score = self.fc(x)
         if isinstance(cls_score, list):
             cls_score = sum(cls_score) / float(len(cls_score))
-        pred = F.softmax(cls_score, dim=1) if cls_score is not None else None
 
-        return self.post_process(pred)
+        if softmax:
+            pred = (
+                F.softmax(cls_score, dim=1) if cls_score is not None else None)
+        else:
+            pred = cls_score
+
+        if post_process:
+            return self.post_process(pred)
+        else:
+            return pred
 
     def forward_train(self, x, gt_label, **kwargs):
-        if isinstance(x, tuple):
-            x = x[-1]
+        x = self.pre_logits(x)
         cls_score = self.fc(x)
         losses = self.loss(cls_score, gt_label, **kwargs)
         return losses

mmcls/models/heads/multi_label_head.py

@@ -1,6 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import torch
-import torch.nn.functional as F
 
 from ..builder import HEADS, build_loss
 from ..utils import is_tracing
@@ -47,14 +46,33 @@ def forward_train(self, cls_score, gt_label, **kwargs):
         losses = self.loss(cls_score, gt_label, **kwargs)
         return losses
 
-    def simple_test(self, x):
+    def pre_logits(self, x):
+        if isinstance(x, tuple):
+            x = x[-1]
+
+        from mmcls.utils import get_root_logger
+        logger = get_root_logger()
+        logger.warning(
+            'The input of MultiLabelClsHead should be already logits. '
+            'Please modify the backbone if you want to get pre-logits feature.'
+        )
+        return x
+
+    def simple_test(self, x, sigmoid=True, post_process=True):
         if isinstance(x, tuple):
             x = x[-1]
         if isinstance(x, list):
             x = sum(x) / float(len(x))
-        pred = F.sigmoid(x) if x is not None else None
 
-        return self.post_process(pred)
+        if sigmoid:
+            pred = torch.sigmoid(x) if x is not None else None
+        else:
+            pred = x
+
+        if post_process:
+            return self.post_process(pred)
+        else:
+            return pred
 
     def post_process(self, pred):
         on_trace = is_tracing()

mmcls/models/heads/multi_label_linear_head.py

@@ -1,6 +1,6 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
 from ..builder import HEADS
 from .multi_label_head import MultiLabelClsHead
@@ -39,21 +39,31 @@ def __init__(self,
 
         self.fc = nn.Linear(self.in_channels, self.num_classes)
 
-    def forward_train(self, x, gt_label, **kwargs):
+    def pre_logits(self, x):
         if isinstance(x, tuple):
             x = x[-1]
+        return x
+
+    def forward_train(self, x, gt_label, **kwargs):
+        x = self.pre_logits(x)
         gt_label = gt_label.type_as(x)
         cls_score = self.fc(x)
         losses = self.loss(cls_score, gt_label, **kwargs)
         return losses
 
-    def simple_test(self, x):
+    def simple_test(self, x, sigmoid=True, post_process=True):
         """Test without augmentation."""
-        if isinstance(x, tuple):
-            x = x[-1]
+        x = self.pre_logits(x)
         cls_score = self.fc(x)
         if isinstance(cls_score, list):
             cls_score = sum(cls_score) / float(len(cls_score))
-        pred = F.sigmoid(cls_score) if cls_score is not None else None
 
-        return self.post_process(pred)
+        if sigmoid:
+            pred = torch.sigmoid(cls_score) if cls_score is not None else None
+        else:
+            pred = cls_score
+
+        if post_process:
+            return self.post_process(pred)
+        else:
+            return pred

mmcls/models/heads/stacked_head.py

@@ -89,9 +89,7 @@ def __init__(self,
         self._init_layers()
 
     def _init_layers(self):
-        self.layers = ModuleList(
-            init_cfg=dict(
-                type='Normal', layer='Linear', mean=0., std=0.01, bias=0.))
+        self.layers = ModuleList()
         in_channels = self.in_channels
         for hidden_channels in self.mid_channels:
             self.layers.append(
@@ -114,24 +112,37 @@ def _init_layers(self):
     def init_weights(self):
         self.layers.init_weights()
 
-    def simple_test(self, x):
-        """Test without augmentation."""
+    def pre_logits(self, x):
         if isinstance(x, tuple):
             x = x[-1]
-        cls_score = x
-        for layer in self.layers:
-            cls_score = layer(cls_score)
+        for layer in self.layers[:-1]:
+            x = layer(x)
+        return x
+
+    @property
+    def fc(self):
+        return self.layers[-1]
+
+    def simple_test(self, x, softmax=True, post_process=True):
+        """Test without augmentation."""
+        x = self.pre_logits(x)
+        cls_score = self.fc(x)
         if isinstance(cls_score, list):
             cls_score = sum(cls_score) / float(len(cls_score))
-        pred = F.softmax(cls_score, dim=1) if cls_score is not None else None
 
-        return self.post_process(pred)
+        if softmax:
+            pred = (
+                F.softmax(cls_score, dim=1) if cls_score is not None else None)
+        else:
+            pred = cls_score
+
+        if post_process:
+            return self.post_process(pred)
+        else:
+            return pred
 
     def forward_train(self, x, gt_label, **kwargs):
-        if isinstance(x, tuple):
-            x = x[-1]
-        cls_score = x
-        for layer in self.layers:
-            cls_score = layer(cls_score)
+        x = self.pre_logits(x)
+        cls_score = self.fc(x)
         losses = self.loss(cls_score, gt_label, **kwargs)
         return losses