Add supcon seg for regression (#2177)

* Fix iseg e2e (#2173)

* Add supcon seg for regression

* Fix det e2e (#2176)

* Fix det e2e

* Fix mypy

* Fix precommit

* Update otx/algorithms/segmentation/configs/ocr_lite_hrnet_18/supcon/__init__.py

* Update otx/algorithms/segmentation/configs/ocr_lite_hrnet_18/supcon/model.py

---------

Co-authored-by: Jaeguk Hyun <[email protected]>

otx/algorithms/segmentation/configs/ocr_lite_hrnet_18/supcon/__init__.py

@@ -0,0 +1,15 @@
+"""Initialization of OCR-Lite-HRnet-18 model for SupCon Segmentation Task."""
+
+# Copyright (C) 2022 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions
+# and limitations under the License.

otx/algorithms/segmentation/configs/ocr_lite_hrnet_18/supcon/data_pipeline.py

@@ -0,0 +1,18 @@
+"""Data Pipeline of HR-Net model for Segmentation Task."""
+
+# Copyright (C) 2023 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions
+# and limitations under the License.
+
+# pylint: disable=invalid-name
+_base_ = ["../../base/data/supcon/data_pipeline.py"]

otx/algorithms/segmentation/configs/ocr_lite_hrnet_18/supcon/hparam.yaml

@@ -0,0 +1,8 @@
+# Hyperparameters.
+hyper_parameters:
+  parameter_overrides:
+    learning_parameters:
+      enable_supcon:
+        default_value: True
+      learning_rate_warmup_iters:
+        default_value: 50

otx/algorithms/segmentation/configs/ocr_lite_hrnet_18/supcon/model.py

@@ -0,0 +1,85 @@
+"""Model configuration of OCR-Lite-HRnet-18 model for SupCon Segmentation Task."""
+
+
+# Copyright (C) 2022 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions
+# and limitations under the License.
+
+# pylint: disable=invalid-name
+
+_base_ = [
+    "../../../../../recipes/stages/segmentation/supcon.py",
+    "../../../../common/adapters/mmcv/configs/backbones/lite_hrnet_18.py",
+]
+
+model = dict(
+    type="SupConDetConB",
+    pretrained=(
+        "https://storage.openvinotoolkit.org/repositories/"
+        "openvino_training_extensions/models/custom_semantic_segmentation/"
+        "litehrnet18_imagenet1k_rsc.pth"
+    ),
+    num_classes=256,
+    num_samples=16,
+    downsample=4,
+    input_transform="resize_concat",
+    in_index=[0, 1, 2, 3],
+    neck=dict(
+        type="SelfSLMLP",
+        in_channels=600,
+        hid_channels=256,
+        out_channels=128,
+        norm_cfg=dict(type="BN1d", requires_grad=True),
+        with_avg_pool=False,
+    ),
+    head=dict(
+        type="DetConHead",
+        predictor=dict(
+            type="SelfSLMLP",
+            in_channels=128,
+            hid_channels=256,
+            out_channels=128,
+            norm_cfg=dict(type="BN1d", requires_grad=True),
+            with_avg_pool=False,
+        ),
+        loss_cfg=dict(type="DetConLoss", temperature=0.1),
+    ),
+    decode_head=dict(
+        type="FCNHead",
+        in_channels=[40, 80, 160, 320],
+        in_index=[0, 1, 2, 3],
+        input_transform="multiple_select",
+        channels=40,
+        kernel_size=1,
+        num_convs=1,
+        concat_input=False,
+        dropout_ratio=-1,
+        num_classes=2,
+        norm_cfg=dict(type="BN", requires_grad=True),
+        align_corners=False,
+        enable_aggregator=True,
+        loss_decode=[
+            dict(
+                type="CrossEntropyLoss",
+                use_sigmoid=False,
+                loss_weight=1.0,
+            ),
+        ],
+    ),
+)
+
+load_from = None
+
+resume_from = None
+
+fp16 = dict(loss_scale=512.0)

tests/regression/regression_config.json

@@ -92,6 +92,11 @@
           "--val-data-roots": "segmentation/regression_voc_cls_decr/val",
           "--test-data-roots": "segmentation/regression_voc_cls_decr/test",
           "--input": "segmentation/regression_voc_cls_decr/test/images"
+        },
+        "supcon": {
+          "--train-data-roots": "segmentation/regression_voc_cls_decr/train",
+          "--val-data-roots": "segmentation/regression_voc_cls_decr/val",
+          "--test-data-roots": "segmentation/regression_voc_cls_decr/test"
         }
       },
       "class_incr": {
@@ -425,6 +430,14 @@
             "Lite-HRNet-18": 0.302,
             "Lite-HRNet-x-mod3": 0.041
           }
+        },
+        "supcon": {
+          "train": {
+            "Lite-HRNet-s-mod2": 0.691,
+            "Lite-HRNet-18-mod2": 0.676,
+            "Lite-HRNet-18": 0.682,
+            "Lite-HRNet-x-mod3": 0.698
+          }
         }
       },
       "class_incr": {
@@ -1601,6 +1614,14 @@
             "Lite-HRNet-18": 446.99,
             "Lite-HRNet-x-mod3": 667.493
           }
+        },
+        "supcon": {
+          "train": {
+            "Lite-HRNet-s-mod2": 241.88,
+            "Lite-HRNet-18-mod2": 279.19,
+            "Lite-HRNet-18": 279.71,
+            "Lite-HRNet-x-mod3": 454.31
+          }
         }
       },
       "class_incr": {
@@ -1993,6 +2014,14 @@
             "Lite-HRNet-18": 43.046,
             "Lite-HRNet-x-mod3": 59.053
           }
+        },
+        "supcon": {
+          "train": {
+            "Lite-HRNet-s-mod2": 13.835,
+            "Lite-HRNet-18-mod2": 16.39,
+            "Lite-HRNet-18": 16.23,
+            "Lite-HRNet-x-mod3": 26.549
+          }
         }
       },
       "class_incr": {

tests/regression/semantic_segmentation/test_segmentation.py

@@ -415,3 +415,74 @@ def test_pot_optimize_eval(self, template, tmp_dir_path):
         result_dict[TASK_TYPE][self.label_type][TRAIN_TYPE]["pot"].append(self.performance)
 
         assert test_result["passed"] is True, test_result["log"]
+
+
+class TestRegressionSupconSegmentation:
+    def setup_method(self):
+        self.label_type = "supcon"
+        self.performance = {}
+
+    def teardown_method(self):
+        with open(f"{result_dir}/result.json", "w") as result_file:
+            json.dump(result_dict, result_file, indent=4)
+
+    @e2e_pytest_component
+    @pytest.mark.parametrize("template", templates, ids=templates_ids)
+    def test_otx_train(self, template, tmp_dir_path):
+        self.performance[template.name] = {}
+
+        tmp_dir_path = tmp_dir_path / "supcon_seg"
+        config_supcon = load_regression_configuration(otx_dir, TASK_TYPE, TRAIN_TYPE, self.label_type)
+        args_supcon = config_supcon["data_path"]
+
+        args_supcon["train_params"] = [
+            "params",
+            "--learning_parameters.num_iters",
+            REGRESSION_TEST_EPOCHS,
+            "--learning_parameters.enable_supcon",
+            "True",
+        ]
+        # Supcon
+        train_start_time = timer()
+        otx_train_testing(template, tmp_dir_path, otx_dir, args_supcon)
+        train_elapsed_time = timer() - train_start_time
+
+        # Evaluation with supcon + supervised training
+        infer_start_time = timer()
+        test_result = regression_eval_testing(
+            template,
+            tmp_dir_path,
+            otx_dir,
+            args_supcon,
+            config_supcon["regression_criteria"]["train"],
+            self.performance[template.name],
+        )
+        infer_elapsed_time = timer() - infer_start_time
+
+        self.performance[template.name][TIME_LOG["train_time"]] = round(train_elapsed_time, 3)
+        self.performance[template.name][TIME_LOG["infer_time"]] = round(infer_elapsed_time, 3)
+        result_dict[TASK_TYPE][self.label_type][TRAIN_TYPE]["train"].append(self.performance)
+
+        assert test_result["passed"] is True, test_result["log"]
+
+    @e2e_pytest_component
+    @pytest.mark.parametrize("template", templates, ids=templates_ids)
+    def test_otx_train_kpi_test(self, template):
+        config_supcon = load_regression_configuration(otx_dir, TASK_TYPE, TRAIN_TYPE, self.label_type)
+        results = result_dict[TASK_TYPE][self.label_type][TRAIN_TYPE]["train"]
+        performance = get_template_performance(results, template)
+
+        kpi_train_result = regression_train_time_testing(
+            train_time_criteria=config_supcon["kpi_e2e_train_time_criteria"]["train"],
+            e2e_train_time=performance[template.name][TIME_LOG["train_time"]],
+            template=template,
+        )
+
+        kpi_eval_result = regression_eval_time_testing(
+            eval_time_criteria=config_supcon["kpi_e2e_eval_time_criteria"]["train"],
+            e2e_eval_time=performance[template.name][TIME_LOG["infer_time"]],
+            template=template,
+        )
+
+        assert kpi_train_result["passed"] is True, kpi_train_result["log"]
+        assert kpi_eval_result["passed"] is True, kpi_eval_result["log"]