Revert h-cls head to linear (#4221)

* Revert h-cls head to linear

* Update changelog

Author: sovrasov

CHANGELOG.md

@@ -72,6 +72,8 @@ All notable changes to this project will be documented in this file.
   (<https://github.com/openvinotoolkit/training_extensions/pull/4199>)
 - Fix label info on loading checkpoint
   (<https://github.com/openvinotoolkit/training_extensions/pull/4200>)
+- Revert h-cls head to linear one
+  (<https://github.com/openvinotoolkit/training_extensions/pull/4221>)
 
 ## \[2.2.2\]
 

src/otx/algo/classification/efficientnet.py

@@ -14,7 +14,7 @@
 from otx.algo.classification.backbones.efficientnet import EFFICIENTNET_VERSION, EfficientNetBackbone
 from otx.algo.classification.classifier import HLabelClassifier, ImageClassifier, SemiSLClassifier
 from otx.algo.classification.heads import (
-    HierarchicalCBAMClsHead,
+    HierarchicalLinearClsHead,
     LinearClsHead,
     MultiLabelLinearClsHead,
     SemiSLLinearClsHead,
@@ -272,11 +272,8 @@ def _build_model(self, head_config: dict) -> nn.Module:
 
         return HLabelClassifier(
             backbone=backbone,
-            neck=nn.Identity(),
-            head=HierarchicalCBAMClsHead(
-                in_channels=backbone.num_features,
-                **copied_head_config,
-            ),
+            neck=GlobalAveragePooling(dim=2),
+            head=HierarchicalLinearClsHead(**copied_head_config, in_channels=backbone.num_features),
             multiclass_loss=nn.CrossEntropyLoss(),
             multilabel_loss=AsymmetricAngularLossWithIgnore(gamma_pos=0.0, gamma_neg=1.0, reduction="sum"),
         )

src/otx/algo/classification/mobilenet_v3.py

@@ -15,7 +15,7 @@
 from otx.algo.classification.backbones import MobileNetV3Backbone
 from otx.algo.classification.classifier import HLabelClassifier, ImageClassifier, SemiSLClassifier
 from otx.algo.classification.heads import (
-    HierarchicalCBAMClsHead,
+    HierarchicalLinearClsHead,
     LinearClsHead,
     MultiLabelNonLinearClsHead,
     SemiSLLinearClsHead,
@@ -314,15 +314,13 @@ def _build_model(self, head_config: dict) -> nn.Module:
 
         copied_head_config = copy(head_config)
         copied_head_config["step_size"] = (ceil(self.input_size[0] / 32), ceil(self.input_size[1] / 32))
-
+        in_channels = MobileNetV3Backbone.MV3_CFG[self.mode]["out_channels"]
         backbone = MobileNetV3Backbone(mode=self.mode, input_size=self.input_size)
+
         return HLabelClassifier(
             backbone=backbone,
-            neck=nn.Identity(),
-            head=HierarchicalCBAMClsHead(
-                in_channels=MobileNetV3Backbone.MV3_CFG[self.mode]["out_channels"],
-                **copied_head_config,
-            ),
+            neck=GlobalAveragePooling(dim=2),
+            head=HierarchicalLinearClsHead(**copied_head_config, in_channels=in_channels),
             multiclass_loss=nn.CrossEntropyLoss(),
             multilabel_loss=AsymmetricAngularLossWithIgnore(gamma_pos=0.0, gamma_neg=1.0, reduction="sum"),
         )

src/otx/algo/classification/timm_model.py

@@ -15,7 +15,7 @@
 from otx.algo.classification.backbones.timm import TimmBackbone
 from otx.algo.classification.classifier import HLabelClassifier, ImageClassifier, SemiSLClassifier
 from otx.algo.classification.heads import (
-    HierarchicalCBAMClsHead,
+    HierarchicalLinearClsHead,
     LinearClsHead,
     MultiLabelLinearClsHead,
     SemiSLLinearClsHead,
@@ -348,11 +348,8 @@ def _build_model(self, head_config: dict) -> nn.Module:
         copied_head_config["step_size"] = (ceil(self.input_size[0] / 32), ceil(self.input_size[1] / 32))
         return HLabelClassifier(
             backbone=backbone,
-            neck=nn.Identity(),
-            head=HierarchicalCBAMClsHead(
-                in_channels=backbone.num_features,
-                **copied_head_config,
-            ),
+            neck=GlobalAveragePooling(dim=2),
+            head=HierarchicalLinearClsHead(**copied_head_config, in_channels=backbone.num_features),
             multiclass_loss=nn.CrossEntropyLoss(),
             multilabel_loss=AsymmetricAngularLossWithIgnore(gamma_pos=0.0, gamma_neg=1.0, reduction="sum"),
         )

src/otx/algo/classification/torchvision_model.py

@@ -14,7 +14,7 @@
 from otx.algo.classification.backbones.torchvision import TorchvisionBackbone, TVModelType
 from otx.algo.classification.classifier import HLabelClassifier, ImageClassifier, SemiSLClassifier
 from otx.algo.classification.heads import (
-    HierarchicalCBAMClsHead,
+    HierarchicalLinearClsHead,
     LinearClsHead,
     MultiLabelLinearClsHead,
     SemiSLLinearClsHead,
@@ -315,11 +315,8 @@ def _build_model(self, head_config: dict) -> nn.Module:
         backbone = TorchvisionBackbone(backbone=self.backbone, pretrained=self.pretrained)
         return HLabelClassifier(
             backbone=backbone,
-            neck=nn.Identity(),
-            head=HierarchicalCBAMClsHead(
-                in_channels=backbone.in_features,
-                **head_config,
-            ),
+            neck=GlobalAveragePooling(dim=2),
+            head=HierarchicalLinearClsHead(**head_config, in_channels=backbone.in_features),
             multiclass_loss=nn.CrossEntropyLoss(),
             multilabel_loss=AsymmetricAngularLossWithIgnore(gamma_pos=0.0, gamma_neg=1.0, reduction="sum"),
         )

src/otx/algo/classification/vit.py

@@ -19,7 +19,7 @@
 from otx.algo.classification.backbones.vision_transformer import VIT_ARCH_TYPE, VisionTransformer
 from otx.algo.classification.classifier import HLabelClassifier, ImageClassifier, SemiSLClassifier
 from otx.algo.classification.heads import (
-    HierarchicalCBAMClsHead,
+    HierarchicalLinearClsHead,
     MultiLabelLinearClsHead,
     SemiSLVisionTransformerClsHead,
     VisionTransformerClsHead,
@@ -466,11 +466,7 @@ def _build_model(self, head_config: dict) -> nn.Module:
         return HLabelClassifier(
             backbone=vit_backbone,
             neck=None,
-            head=HierarchicalCBAMClsHead(
-                in_channels=vit_backbone.embed_dim,
-                step_size=1,
-                **head_config,
-            ),
+            head=HierarchicalLinearClsHead(**head_config, in_channels=vit_backbone.embed_dim),
             multiclass_loss=nn.CrossEntropyLoss(),
             multilabel_loss=AsymmetricAngularLossWithIgnore(gamma_pos=0.0, gamma_neg=1.0, reduction="sum"),
             init_cfg=init_cfg,