feature: export head for moby

configs/selfsup/moby/moby_deit_small_4xb32_100e_jpg.py

@@ -88,5 +88,5 @@
 
 # export config
 # export = dict(export_neck=True)
-export = dict(export_neck=False)
+export = dict(export_neck=False, export_head=False)
 checkpoint_sync_export = True

configs/selfsup/moby/moby_deit_small_p16_4xb128_300e_tfrecord.py

@@ -104,5 +104,5 @@
 total_epochs = 300
 
 # export config
-export = dict(export_neck=False)
+export = dict(export_neck=False, export_head=False)
 checkpoint_sync_export = True

configs/selfsup/moby/moby_dynamic_swin_tiny_8xb64_300e_tfrecord.py

@@ -91,5 +91,5 @@
 total_epochs = 300
 
 # export config
-export = dict(export_neck=False)
+export = dict(export_neck=False, export_head=False)
 checkpoint_sync_export = True

configs/selfsup/moby/moby_rn50_4xb128_100e_tfrecord.py

@@ -111,5 +111,5 @@
 total_epochs = 100
 
 # export config
-export = dict(export_neck=True)
+export = dict(export_neck=True, export_head=True)
 checkpoint_sync_export = True

easycv/apis/export.py

@@ -468,7 +468,15 @@ def _export_moco(model, cfg, filename):
 
 
 def _export_moby(model, cfg, filename):
-    """ export model and preprocess config
+    """ export model and preprocess config:
+        can export backbone only,
+        or export backbone and neck (projector_q),
+        or export backbone, neck (projector_q) and head (predictor)
+
+        Modified by YANG Ruixin
+        GitHub: https://github.com/yang-ruixin
+        Email: yang_ruixin@126.com
+        Date: 2023/01/05
 
     Args:
         model (nn.Module):  model to be exported
@@ -478,8 +486,9 @@ def _export_moby(model, cfg, filename):
     if hasattr(cfg, 'export'):
         export_cfg = cfg.export
     else:
-        export_cfg = dict(export_neck=False)
+        export_cfg = dict(export_neck=False, export_head=False)
     export_neck = export_cfg.get('export_neck', False)
+    export_head = export_cfg.get('export_head', False)
 
     model_config = dict(
         type='Classification',
@@ -492,7 +501,11 @@ def _export_moby(model, cfg, filename):
             num_classes=1000,
         ),
     )
-    if export_neck:
+
+    if export_head:
+        model_config['neck'] = cfg.model.neck
+        model_config['head'] = cfg.model.head
+    elif export_neck:
         model_config['neck'] = cfg.model.neck
 
     img_norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
@@ -510,11 +523,20 @@ def _export_moby(model, cfg, filename):
     meta = dict(config=json.dumps(config))
 
     state_dict = OrderedDict()
+    neck_key = 'projector_q'
+    head_key = 'predictor'
     for k, v in model.state_dict().items():
         if k.startswith('backbone'):
             state_dict[k] = v
-        neck_key = 'projector_q'
-        if export_neck and k.startswith(neck_key):
+
+        if export_head:
+            if k.startswith(neck_key):
+                new_key = k.replace(neck_key, 'neck_0')
+                state_dict[new_key] = v
+            elif k.startswith(head_key):
+                new_key = k.replace(head_key, 'head_0')
+                state_dict[new_key] = v
+        elif export_neck and k.startswith(neck_key):
             new_key = k.replace(neck_key, 'neck_0')
             state_dict[new_key] = v
 

easycv/models/classification/classification.py

@@ -241,25 +241,53 @@ def forward_feature(self, img) -> Dict[str, torch.Tensor]:
         """Forward feature  means forward backbone  + neck/multineck ,get dict of output feature,
             self.neck_num = 0: means only forward backbone, output backbone feature with avgpool, with key neck,
             self.neck_num > 0: means has 1/multi neck, output neck's feature with key neck_neckidx_featureidx, suck as neck_0_0
+
+            Can also forward head (predictor) when export_head=True in the config file
+            for self-supervised learning MoBY algorithm (one neck + one head),
+            head feature could be used for searching images by image
+
+            Modified by YANG Ruixin
+            GitHub: https://github.com/yang-ruixin
+            Email: yang_ruixin@126.com
+            Date: 2023/01/05
+
         Returns:
             x (torch.Tensor): feature tensor
         """
         return_dict = {}
         x = self.backbone(img)
         # return_dict['backbone'] = x[-1]
+
         if hasattr(self, 'neck_0'):
-            tmp = []
+            tmp_neck = []
             for idx in range(self.neck_num):
                 neck_name = 'neck_%d' % idx
                 h = getattr(self, neck_name)
                 neck_h = h([i for i in x])
-                tmp = tmp + neck_h
+                tmp_neck = tmp_neck + neck_h
                 for j in range(len(neck_h)):
                     neck_name = 'neck_%d_%d' % (idx, j)
                     return_dict['neck_%d_%d' % (idx, j)] = neck_h[j]
                     if neck_name not in self.extract_list:
                         self.extract_list.append(neck_name)
-            return_dict['neck'] = tmp[0]
+
+            if (hasattr(self, 'head_0')
+                    and 'MoBYMLP' in str(self.__dict__['_modules']['head_0'])
+                    and 'MoBYMLP' in str(self.__dict__['_modules']['neck_0'])):
+                tmp_head = []
+                for idx in range(self.head_num):
+                    head_name = 'head_%d' % idx
+                    h = getattr(self, head_name)
+                    head_h = h([i for i in tmp_neck])
+                    tmp_head = tmp_head + head_h
+                    for j in range(len(head_h)):
+                        head_name = 'head_%d_%d' % (idx, j)
+                        return_dict['head_%d_%d' % (idx, j)] = head_h[j]
+                        if head_name not in self.extract_list:
+                            self.extract_list.append(head_name)
+                return_dict['neck'] = tmp_head[0]  # actually we return return_dict['head']
+            else:
+                return_dict['neck'] = tmp_neck[0]
         else:
             feature = self.avg_pool(x[-1])
             feature = feature.view(feature.size(0), -1)

easycv/models/heads/__init__.py

@@ -4,3 +4,4 @@
 from .latent_pred_head import LatentPredictHead
 from .mp_metric_head import MpMetrixHead
 from .multi_cls_head import MultiClsHead
+from .moby_head import MoBYMLP

easycv/models/heads/moby_head.py

@@ -0,0 +1,62 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+#
+# Used for self-supervised learning MoBY algorithm, when export_head=True in the config file
+#
+# Author: YANG Ruixin
+# GitHub: https://github.com/yang-ruixin
+# Email: yang_ruixin@126.com
+# Date: 2023/01/05
+
+# from typing import Dict, List
+
+# import torch
+import torch.nn as nn
+# from mmcv.cnn.utils.weight_init import initialize
+#
+# from easycv.core.evaluation.metrics import accuracy
+# from easycv.utils.logger import get_root_logger
+# from easycv.utils.registry import build_from_cfg
+from ..registry import HEADS  # , LOSSES
+
+from ..utils import _init_weights
+
+
+@HEADS.register_module
+class MoBYMLP(nn.Module):
+
+    def __init__(self,
+                 in_channels=256,
+                 hid_channels=4096,
+                 out_channels=256,
+                 num_layers=2,
+                 with_avg_pool=True):
+        super(MoBYMLP, self).__init__()
+
+        # hidden layers
+        linear_hidden = [nn.Identity()]
+        for i in range(num_layers - 1):
+            linear_hidden.append(
+                nn.Linear(in_channels if i == 0 else hid_channels,
+                          hid_channels))
+            linear_hidden.append(nn.BatchNorm1d(hid_channels))
+            linear_hidden.append(nn.ReLU(inplace=True))
+        self.linear_hidden = nn.Sequential(*linear_hidden)
+        self.linear_out = nn.Linear(
+            in_channels if num_layers == 1 else hid_channels,
+            out_channels) if num_layers >= 1 else nn.Identity()
+        self.with_avg_pool = True
+        self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))
+
+    def forward(self, x):
+        x = x[0]
+        if self.with_avg_pool and len(x.shape) == 4:
+            bs = x.shape[0]
+            x = self.avg_pool(x).view([bs, -1])
+        # print(x.shape)
+        # exit()
+        x = self.linear_hidden(x)
+        x = self.linear_out(x)
+        return [x]
+
+    def init_weights(self, init_linear='normal'):
+        _init_weights(self, init_linear)