From c479e4eecce2f642d7efd3dd9d90d97b4ecd7114 Mon Sep 17 00:00:00 2001 From: yang-ruixin Date: Sun, 8 Jan 2023 23:40:01 +0800 Subject: [PATCH] feature: export head for moby --- .../moby/moby_deit_small_4xb32_100e_jpg.py | 2 +- ...oby_deit_small_p16_4xb128_300e_tfrecord.py | 2 +- ...y_dynamic_swin_tiny_8xb64_300e_tfrecord.py | 2 +- .../moby/moby_rn50_4xb128_100e_tfrecord.py | 2 +- easycv/apis/export.py | 32 ++++++++-- .../models/classification/classification.py | 34 +++++++++- easycv/models/heads/__init__.py | 1 + easycv/models/heads/moby_head.py | 62 +++++++++++++++++++ 8 files changed, 125 insertions(+), 12 deletions(-) create mode 100644 easycv/models/heads/moby_head.py diff --git a/configs/selfsup/moby/moby_deit_small_4xb32_100e_jpg.py b/configs/selfsup/moby/moby_deit_small_4xb32_100e_jpg.py index fbbd88a8..b7988e0e 100644 --- a/configs/selfsup/moby/moby_deit_small_4xb32_100e_jpg.py +++ b/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 diff --git a/configs/selfsup/moby/moby_deit_small_p16_4xb128_300e_tfrecord.py b/configs/selfsup/moby/moby_deit_small_p16_4xb128_300e_tfrecord.py index b878a026..206974a3 100644 --- a/configs/selfsup/moby/moby_deit_small_p16_4xb128_300e_tfrecord.py +++ b/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 diff --git a/configs/selfsup/moby/moby_dynamic_swin_tiny_8xb64_300e_tfrecord.py b/configs/selfsup/moby/moby_dynamic_swin_tiny_8xb64_300e_tfrecord.py index 18c9ca58..de29eadd 100644 --- a/configs/selfsup/moby/moby_dynamic_swin_tiny_8xb64_300e_tfrecord.py +++ b/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 diff --git a/configs/selfsup/moby/moby_rn50_4xb128_100e_tfrecord.py b/configs/selfsup/moby/moby_rn50_4xb128_100e_tfrecord.py index 7e99c17d..375b9308 100644 --- a/configs/selfsup/moby/moby_rn50_4xb128_100e_tfrecord.py +++ b/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 diff --git a/easycv/apis/export.py b/easycv/apis/export.py index 9075c61a..4cee03ff 100644 --- a/easycv/apis/export.py +++ b/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 diff --git a/easycv/models/classification/classification.py b/easycv/models/classification/classification.py index 57c32c08..6bb6eb6f 100644 --- a/easycv/models/classification/classification.py +++ b/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) diff --git a/easycv/models/heads/__init__.py b/easycv/models/heads/__init__.py index cf54aa35..d8a4225c 100644 --- a/easycv/models/heads/__init__.py +++ b/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 diff --git a/easycv/models/heads/moby_head.py b/easycv/models/heads/moby_head.py new file mode 100644 index 00000000..92ed3101 --- /dev/null +++ b/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)