Add EfficientNetV2 (#9)

* Apply MPA public repo changes to inner source (#8)

* Move tasks & model templates to OTE repo
* Update external/training_extension to MPA PR
* Remove external/mmdetection
* Update init_venv.sh
* Update README.md
* Enable non-incremental learning
* TODO: Remove OTE SDK/Task dependency

* Update mpa's minor issue and add model (#10)

* Add det/seg OTELoggerHook & OTEProgressHook

* Add effnetv2 template & seg export config update

* Add ote hooks

* Add metrics score to best_acc

* Fix some pretrained var conf

* Separate mpa changes and apis task changes

* Update SamClassifier & seg weight mixing

* Minor fix

* Update segmentation configuration

* Update sam_classifier.py

* Merge remote-tracking branch 'origin/ote' into ote-public

Co-authored-by: Harim Kang <[email protected]>
Co-authored-by: Jihwan Eom <[email protected]>

Author: 3 people

README.md

@@ -44,7 +44,10 @@ related NN *model* architecures.
 ### TL Tasks
 * Classification
 * Detection
-* Semantic segmentation
+* Segmentation
+    * (WIP) Instance segmentation
+    * Semantic segmentation
+    * (TBD) Panoptic segmentation
 > Train / Infer / Evaluate / Export operations are supported for each task
 
 ### TL Methods

init_venv.sh

@@ -67,7 +67,7 @@ if [ -e "$CUDA_HOME" ]; then
   fi
 fi
 
-# install PyTorch and MMCV.
+# Install PyTorch and MMCV.
 export TORCH_VERSION=1.8.2
 export TORCHVISION_VERSION=0.9.2
 export MMCV_VERSION=1.3.14
@@ -109,10 +109,10 @@ pip install torch==${TORCH_VERSION} torchvision==${TORCHVISION_VERSION} -f https
 
 # Install mmcv
 pip install --no-cache-dir mmcv-full==${MMCV_VERSION} -c ${CONSTRAINTS_FILE} || exit 1
+sed -i "s/force=False/force=True/g" ${venv_dir}/lib/python${PYTHON_VERSION}/site-packages/mmcv/utils/registry.py  # Patch: remedy for MMCV registry collision from mmdet/mmseg
 
-# Install other requirements.
 # Install mmpycocotools from source to make sure it is compatible with installed numpy version.
-pip install --no-cache-dir --no-binary=mmpycocotools mmpycocotools || exit 1
+pip install --no-cache-dir --no-binary=mmpycocotools mmpycocotools -c ${CONSTRAINTS_FILE} || exit 1
 cat requirements.txt | xargs -n 1 -L 1 pip install --no-cache || exit 1
 
 # Install external modules
@@ -134,7 +134,7 @@ else
 fi
 
 # Install MPA
-pip install -e . || exit 1
+pip install -e . -c ${CONSTRAINTS_FILE} || exit 1
 MPA_DIR=`realpath .`
 echo "export MPA_DIR=${MPA_DIR}" >> ${venv_dir}/bin/activate
 

models/classification/_base_/ote_efficientnet_v2.py

@@ -0,0 +1,13 @@
+# model settings
+model = dict(
+    type='ImageClassifier',
+    backbone=dict(
+        type='OTEEfficientNetV2',
+        version='s_21k'),
+    neck=dict(type='GlobalAveragePooling'),
+    head=dict(
+        type='LinearClsHead',
+        num_classes=1000,
+        in_channels=1280,
+        loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
+    ))

models/classification/ote_efficientnet_v2_s.yaml

@@ -0,0 +1,7 @@
+_base_: _base_/ote_efficientnet_v2.py
+
+model:
+    type: SAMImageClassifier
+    task: classification
+    backbone:
+        version: s_21k

mpa/cls/exporter.py

@@ -78,7 +78,7 @@ def run(self, model_cfg, model_ckpt, data_cfg, **kwargs):
                               input_names=['data'],
                               output_names=['logits', 'features', 'vector'],
                               dynamic_axes={},
-                              opset_version=9,
+                              opset_version=11,
                               operator_export_type=torch.onnx.OperatorExportTypes.ONNX
                               )
 

mpa/cls/stage.py

@@ -3,7 +3,6 @@
 import numpy as np
 
 from mmcv import ConfigDict
-from mmcv import Config
 from mmcv import build_from_cfg
 
 from mpa.stage import Stage
@@ -24,12 +23,7 @@ def configure(self, model_cfg, model_ckpt, data_cfg, training=True, **kwargs):
         logger.info(f'configure: training={training}')
 
         # Recipe + model
-        cfg = Config(
-            self.cfg._cfg_dict,
-            self.cfg.text,
-            self.cfg.filename
-        )
-
+        cfg = self.cfg
         if model_cfg:
             if hasattr(cfg, 'model'):
                 cfg.merge_from_dict(model_cfg._cfg_dict)
@@ -54,9 +48,9 @@ def configure(self, model_cfg, model_ckpt, data_cfg, training=True, **kwargs):
             cfg.model.backbone.model_path = ir_path
 
         pretrained = kwargs.get('pretrained', None)
-        if pretrained:
+        if pretrained and isinstance(pretrained, str):
             logger.info(f'Overriding cfg.load_from -> {pretrained}')
-            cfg.load_from = pretrained  # Overriding by stage input
+            cfg.load_from = pretrained
 
         # Data
         if data_cfg:

mpa/det/inferrer.py

@@ -127,8 +127,20 @@ def infer(self, cfg):
         # Inference
         model = MMDataParallel(model, device_ids=[0])
         detections = single_gpu_test(model, data_loader)
+
+        eval_cfg = cfg.evaluation.copy()
+        eval_cfg.pop('interval', None)
+        eval_cfg.pop('save_best', None)
+
+        metric = dataset.evaluate(
+            detections,
+            logger='silent',
+            **eval_cfg
+        )[cfg.evaluation.metric]
+
         outputs = dict(
             classes=target_classes,
             detections=detections,
+            metric=metric,
         )
         return outputs

mpa/det/stage.py

@@ -37,7 +37,7 @@ def configure(self, model_cfg, model_ckpt, data_cfg, training=True, **kwargs):
         if model_ckpt:
             cfg.load_from = self.get_model_ckpt(model_ckpt)
         pretrained = kwargs.get('pretrained', None)
-        if pretrained:
+        if pretrained and isinstance(pretrained, str):
             logger.info(f'Overriding cfg.load_from -> {pretrained}')
             cfg.load_from = pretrained  # Overriding by stage input
 

mpa/det/trainer.py

@@ -90,7 +90,7 @@ def run(self, model_cfg, model_ckpt, data_cfg, **kwargs):
         # Metadata
         meta = dict()
         meta['env_info'] = env_info
-        meta['config'] = cfg.pretty_text
+        # meta['config'] = cfg.pretty_text
         meta['seed'] = cfg.seed
         meta['exp_name'] = cfg.work_dir
         if cfg.checkpoint_config is not None:
@@ -108,8 +108,8 @@ def run(self, model_cfg, model_ckpt, data_cfg, **kwargs):
                 cfg.optimizer.lr = new_lr
 
         # Save config
-        cfg.dump(osp.join(cfg.work_dir, 'config.py'))
-        logger.info(f'Config:\n{cfg.pretty_text}')
+        # cfg.dump(osp.join(cfg.work_dir, 'config.py'))
+        # logger.info(f'Config:\n{cfg.pretty_text}')
 
         if distributed:
             os.environ['MASTER_ADDR'] = cfg.dist_params.get('master_addr', 'localhost')

mpa/modules/datasets/mpa_det_dataset.py

@@ -1,5 +1,5 @@
 from mmdet.datasets.builder import DATASETS
-from detection_tasks.extension.datasets import OTEDataset
+from detection_tasks.extension.datasets import OTEDataset  # TODO: Remove OTE SDK/Task dependency
 from mpa.utils.logger import get_logger
 
 logger = get_logger()

mpa/modules/datasets/mpa_seg_incr_dataset.py

@@ -1,5 +1,5 @@
 from mmseg.datasets.builder import DATASETS
-from segmentation_tasks.extension.datasets import OTEDataset
+from segmentation_tasks.extension.datasets import OTEDataset  # TODO: Remove OTE SDK/Task dependency
 from mpa.utils.logger import get_logger
 
 logger = get_logger()

mpa/modules/datasets/pipelines/mpa_cls_pipeline.py

@@ -1,5 +1,5 @@
 from mmcls.datasets import PIPELINES
-from detection_tasks.extension.utils import LoadImageFromOTEDataset
+from detection_tasks.extension.utils import LoadImageFromOTEDataset  # TODO: Remove OTE SDK/Task dependency
 
 
 @PIPELINES.register_module()

mpa/modules/hooks/early_stopping_hook.py

@@ -1,5 +1,5 @@
 from mmcv.runner.hooks import HOOKS
-from detection_tasks.extension.utils.hooks import EarlyStoppingHook
+from detection_tasks.extension.utils.hooks import EarlyStoppingHook  # TODO: Remove OTE SDK/Task dependency
 
 
 @HOOKS.register_module()

mpa/modules/models/backbones/__init__.py

@@ -1,4 +1,5 @@
 # flake8: noqa
 from . import wideresnet
 from . import mobilenetv3
-from . import efficientnet
+from . import efficientnet
+from . import efficientnetv2

mpa/modules/models/backbones/efficientnetv2.py

@@ -0,0 +1,99 @@
+"""
+    EfficientNet for ImageNet-1K, implemented in PyTorch.
+    Original papers:
+    - 'EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks,' https://arxiv.org/abs/1905.11946,
+    - 'Adversarial Examples Improve Image Recognition,' https://arxiv.org/abs/1911.09665.
+"""
+
+import os
+
+import timm
+import torch.nn as nn
+from mmcls.models.builder import BACKBONES
+from mmcv.runner import load_checkpoint
+from mpa.utils.logger import get_logger
+
+logger = get_logger()
+
+pretrained_root = "https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-effv2-weights/"
+pretrained_urls = {
+    "efficientnetv2_s_21k": pretrained_root + "tf_efficientnetv2_s_21k-6337ad01.pth",
+    "efficientnetv2_s_1k": pretrained_root + "tf_efficientnetv2_s_21ft1k-d7dafa41.pth",
+}
+
+NAME_DICT = {
+                'mobilenetv3_large_21k': 'mobilenetv3_large_100_miil_in21k',
+                'mobilenetv3_large_1k': 'mobilenetv3_large_100_miil',
+                'tresnet': 'tresnet_m',
+                'efficientnetv2_s_21k': 'tf_efficientnetv2_s_in21k',
+                'efficientnetv2_s_1k': 'tf_efficientnetv2_s_in21ft1k',
+                'efficientnetv2_m_21k': 'tf_efficientnetv2_m_in21k',
+                'efficientnetv2_m_1k': 'tf_efficientnetv2_m_in21ft1k',
+                'efficientnetv2_b0': 'tf_efficientnetv2_b0',
+            }
+
+
+class TimmModelsWrapper(nn.Module):
+    def __init__(self,
+                 model_name,
+                 pretrained=True,
+                 pooling_type='avg',
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.model_name = model_name
+        self.pretrained = pretrained
+        self.is_mobilenet = True if model_name in [
+                "mobilenetv3_large_100_miil_in21k", "mobilenetv3_large_100_miil"
+            ] else False
+        self.model = timm.create_model(NAME_DICT[self.model_name],
+                                       pretrained=pretrained,
+                                       num_classes=1000)
+        self.model.classifier = None  # Detach classifier. Only use 'backbone' part in mpa.
+        self.num_head_features = self.model.num_features
+        self.num_features = (self.model.conv_head.in_channels if self.is_mobilenet
+                             else self.model.num_features)
+        self.pooling_type = pooling_type
+
+    def forward(self, x, return_featuremaps=True, **kwargs):
+        y = self.extract_features(x)
+        if return_featuremaps:
+            return y
+
+    def extract_features(self, x):
+        if self.is_mobilenet:
+            x = self.model.conv_stem(x)
+            x = self.model.bn1(x)
+            x = self.model.act1(x)
+            y = self.model.blocks(x)
+            return y
+        return self.model.forward_features(x)
+
+    def get_config_optim(self, lrs):
+        parameters = [
+            {'params': self.model.named_parameters()},
+        ]
+        if isinstance(lrs, list):
+            assert len(lrs) == len(parameters)
+            for lr, param_dict in zip(lrs, parameters):
+                param_dict['lr'] = lr
+        else:
+            assert isinstance(lrs, float)
+            for param_dict in parameters:
+                param_dict['lr'] = lrs
+
+        return parameters
+
+
+@BACKBONES.register_module()
+class OTEEfficientNetV2(TimmModelsWrapper):
+    def __init__(self, version="s_21k", **kwargs):
+        self.model_name = "efficientnetv2_" + version
+        super().__init__(model_name=self.model_name, **kwargs)
+
+    def init_weights(self, pretrained=None):
+        if isinstance(pretrained, str) and os.path.exists(pretrained):
+            load_checkpoint(self, pretrained)
+            logger.info(f"init weight - {pretrained}")
+        elif pretrained is not None:
+            load_checkpoint(self, pretrained_urls[self.model_name])
+            logger.info(f"init weight - {pretrained_urls[self.model_name]}")