[HOTFIX][RELEASE1.0] Fixed label_scheme mismatch in classification (#1841)

* Fix some cls-patch

* Remove comments

* Fix some var

Author: harimkang

otx/mpa/cls/exporter.py

@@ -41,7 +41,6 @@ def model_builder_helper(*args, **kwargs):
             return model
 
         kwargs["model_builder"] = model_builder_helper
-
         return super().run(model_cfg, model_ckpt, data_cfg, **kwargs)
 
     @staticmethod

otx/mpa/cls/incremental/stage.py

@@ -4,7 +4,7 @@
 
 from mmcv import ConfigDict
 
-from otx.mpa.cls.stage import ClsStage, Stage
+from otx.mpa.cls.stage import ClsStage
 from otx.mpa.utils.config_utils import update_or_add_custom_hook
 from otx.mpa.utils.logger import get_logger
 
@@ -35,9 +35,7 @@ def configure_task(self, cfg, training, **kwargs):
     # noqa: C901
     def configure_task_adapt(self, cfg, training, **kwargs):
         """Configure for Task Adaptation Task"""
-
-        self.adapt_type = cfg["task_adapt"].get("op", "REPLACE")
-        train_data_cfg = Stage.get_data_cfg(cfg, "train")
+        train_data_cfg = self.get_data_cfg(cfg, "train")
         if training:
             if train_data_cfg.type not in CLASS_INC_DATASET:
                 logger.warning(f"Class Incremental Learning for {train_data_cfg.type} is not yet supported!")
@@ -46,24 +44,17 @@ def configure_task_adapt(self, cfg, training, **kwargs):
 
             if cfg.model.type in WEIGHT_MIX_CLASSIFIER:
                 cfg.model.task_adapt = ConfigDict(
-                    src_classes=self.model_classes,
-                    dst_classes=self.data_classes,
+                    src_classes=self.org_model_classes,
+                    dst_classes=self.model_classes,
                 )
             else:
                 logger.warning(f"Weight mixing for {cfg.model.type} is not yet supported!")
 
-            # refine self.dst_class following adapt_type (REPLACE, MERGE)
-            self.refine_classes(train_data_cfg)
-            cfg.model.head.num_classes = len(self.dst_classes)
+            train_data_cfg.classes = self.model_classes
 
             # configure loss, sampler, task_adapt_hook
             self.configure_task_modules(cfg)
 
-        else:  # if eval phase (eval)
-            if train_data_cfg.get("new_classes"):
-                self.refine_classes(train_data_cfg)
-                cfg.model.head.num_classes = len(self.dst_classes)
-
     def configure_task_modules(self, cfg):
         if not cfg.model.get("multilabel", False) and not cfg.model.get("hierarchical", False):
             efficient_mode = cfg["task_adapt"].get("efficient_mode", True)
@@ -73,8 +64,8 @@ def configure_task_modules(self, cfg):
             efficient_mode = cfg["task_adapt"].get("efficient_mode", False)
             sampler_type = "cls_incr"
 
-        if len(set(self.model_classes) & set(self.dst_classes)) == 0 or set(self.model_classes) == set(
-            self.dst_classes
+        if len(set(self.org_model_classes) & set(self.model_classes)) == 0 or set(self.org_model_classes) == set(
+            self.model_classes
         ):
             sampler_flag = False
         else:
@@ -83,8 +74,8 @@ def configure_task_modules(self, cfg):
         # Update Task Adapt Hook
         task_adapt_hook = ConfigDict(
             type="TaskAdaptHook",
-            src_classes=self.old_classes,
-            dst_classes=self.dst_classes,
+            src_classes=self.org_model_classes,
+            dst_classes=self.model_classes,
             model_type=cfg.model.type,
             sampler_flag=sampler_flag,
             sampler_type=sampler_type,
@@ -93,8 +84,8 @@ def configure_task_modules(self, cfg):
         update_or_add_custom_hook(cfg, task_adapt_hook)
 
     def configure_loss(self, cfg):
-        if len(set(self.model_classes) & set(self.dst_classes)) == 0 or set(self.model_classes) == set(
-            self.dst_classes
+        if len(set(self.org_model_classes) & set(self.model_classes)) == 0 or set(self.org_model_classes) == set(
+            self.model_classes
         ):
             cfg.model.head.loss = dict(type="CrossEntropyLoss", loss_weight=1.0)
         else:
@@ -104,20 +95,6 @@ def configure_loss(self, cfg):
             )
             ib_loss_hook = ConfigDict(
                 type="IBLossHook",
-                dst_classes=self.dst_classes,
+                dst_classes=self.model_classes,
             )
             update_or_add_custom_hook(cfg, ib_loss_hook)
-
-    def refine_classes(self, train_cfg):
-        # Get 'new_classes' in data.train_cfg & get 'old_classes' pretreained model meta data CLASSES
-        new_classes = train_cfg["new_classes"]
-        self.old_classes = self.model_meta["CLASSES"]
-        if self.adapt_type == "REPLACE":
-            # if 'REPLACE' operation, then self.dst_classes -> data_classes
-            self.dst_classes = self.data_classes.copy()
-        elif self.adapt_type == "MERGE":
-            # if 'MERGE' operation, then self.dst_classes -> old_classes + new_classes (merge)
-            self.dst_classes = self.old_classes + [cls for cls in new_classes if cls not in self.old_classes]
-        else:
-            raise KeyError(f"{self.adapt_type} is not supported for task_adapt options!")
-        train_cfg.classes = self.dst_classes

otx/mpa/cls/stage.py

@@ -143,25 +143,27 @@ def configure_task(self, cfg, training, **kwargs):
         self.configure_classes(cfg)
 
     def configure_classes(self, cfg):
-        model_classes, data_classes = [], []
-        self.model_meta = self.get_model_meta(cfg)
-        train_data_cfg = Stage.get_data_cfg(cfg, "train")
-        if isinstance(train_data_cfg, list):
-            train_data_cfg = train_data_cfg[0]
-
-        model_classes = Stage.get_model_classes(cfg)
-        data_classes = Stage.get_data_classes(cfg)
-
-        if cfg.get("model_classes", []):
-            cfg.model.head.num_classes = len(cfg.model_classes)  # read from prev model label_schema
-        elif model_classes:
-            cfg.model.head.num_classes = len(model_classes)  # read from ckpt meta
-        elif data_classes:
-            cfg.model.head.num_classes = len(data_classes)  # read from env label_schema
-        self.model_meta["CLASSES"] = model_classes
-
-        if not train_data_cfg.get("new_classes", False):  # when train_data_cfg doesn't have 'new_classes' key
-            new_classes = np.setdiff1d(data_classes, model_classes).tolist()
-            train_data_cfg["new_classes"] = new_classes
+        """Patch classes for model and dataset."""
+        self.task_adapt_op = "REPLACE"
+        if "task_adapt" in cfg:
+            self.task_adapt_op = cfg["task_adapt"].get("op", "REPLACE")
+
+        org_model_classes = self.get_model_classes(cfg)
+        data_classes = self.get_data_classes(cfg)
+
+        # Model classes
+        if self.task_adapt_op == "REPLACE":
+            if len(data_classes) == 0:
+                model_classes = org_model_classes.copy()
+            else:
+                model_classes = data_classes.copy()
+        elif self.task_adapt_op == "MERGE":
+            model_classes = org_model_classes + [cls for cls in data_classes if cls not in org_model_classes]
+        else:
+            raise KeyError(f"{self.task_adapt_op} is not supported for task_adapt options!")
+
+        # Model architecture
+        cfg.model.head.num_classes = len(model_classes)
+
+        self.org_model_classes = org_model_classes
         self.model_classes = model_classes
-        self.data_classes = data_classes

otx/mpa/cls/trainer.py

@@ -64,13 +64,6 @@ def run(self, model_cfg, model_ckpt, data_cfg, **kwargs):  # noqa: C901
                 cfg.checkpoint_config.meta["tasks"] = repr_ds.tasks
             if hasattr(repr_ds, "CLASSES"):
                 cfg.checkpoint_config.meta["CLASSES"] = repr_ds.CLASSES
-            if "task_adapt" in cfg:
-                if hasattr(self, "model_tasks"):  # for incremnetal learning
-                    cfg.checkpoint_config.meta.update({"tasks": self.model_tasks})
-                    # instead of update(self.old_tasks), update using "self.model_tasks"
-                if self.model_classes:
-                    cfg.checkpoint_config.meta.update({"CLASSES": self.model_classes})
-                    # FIXME:self.dst_classes?
 
         self.configure_samples_per_gpu(cfg, "train", self.distributed)
         self.configure_fp16_optimizer(cfg, self.distributed)
@@ -85,10 +78,6 @@ def run(self, model_cfg, model_ckpt, data_cfg, **kwargs):  # noqa: C901
 
         self.configure_compat_cfg(cfg)
 
-        # Save config
-        # cfg.dump(osp.join(cfg.work_dir, 'config.py'))
-        # logger.info(f'Config:\n{cfg.pretty_text}')
-
         # register custom eval hooks
         validate = True if cfg.data.get("val", None) else False
         if validate:

tests/unit/mpa/cls/incremental/test_cls_incremental_stage.py

@@ -15,28 +15,28 @@ def setup(self) -> None:
 
     @pytest.mark.parametrize("mode", ["MERGE", "REPLACE"])
     @e2e_pytest_unit
-    def test_configure_task_adapt(self, mode):
+    def test_configure_classes(self, mode, mocker):
+
         self.stage.cfg.merge_from_dict(self.data_cfg)
         self.stage.cfg.task_adapt.op = mode
-        model_classes = ["label_0", "label_1", "label_3", "label_n"]
-        self.stage.model_meta = {"CLASSES": model_classes}
-        self.stage.model_classes = model_classes
-        self.stage.data_classes = self.data_cfg.data.train.data_classes
-        self.stage.new_classes = self.data_cfg.data.train.data_classes
-        self.stage.configure_task_adapt(self.stage.cfg, True)
+        origin_model_classes = ["label_0", "label_3", "label_n"]
+        mocker.patch("otx.mpa.cls.incremental.stage.IncrClsStage.get_model_classes", return_value=origin_model_classes)
+        self.stage.data_classes = self.data_cfg.data.train.data_classes  # ["label_0", "label_1"]
+        self.stage.configure_classes(self.stage.cfg)
+        merge_target = ["label_0", "label_1", "label_3", "label_n"]
 
         if mode == "REPLACE":
             assert self.stage.cfg.model.head.num_classes == len(self.stage.data_classes)
         else:
-            assert self.stage.cfg.model.head.num_classes == len(self.stage.data_classes) + len(self.stage.new_classes)
+            assert self.stage.cfg.model.head.num_classes == len(merge_target)
 
     @e2e_pytest_unit
     def test_configure_task_modules(self, monkeypatch, mocker):
         mock_update_hook = mocker.patch("otx.mpa.cls.incremental.stage.update_or_add_custom_hook")
         # some dummy classes
+        self.stage.data_classes = [0, 1]
         self.stage.model_classes = [0, 1]
-        self.stage.dst_classes = [0, 1]
-        self.stage.old_classes = [0, 1]
+        self.stage.org_model_classes = [0, 1]
         self.stage.configure_task_modules(self.stage.cfg)
 
         mock_update_hook.assert_called_once()