Fix H-label & export with explain

openvinotoolkit · Apr 5, 2024 · 249035b · 249035b
1 parent d4c71e9
commit 249035b
Show file tree

Hide file tree

Showing 8 changed files with 203 additions and 503 deletions.
diff --git a/src/otx/algo/classification/backbones/__init__.py b/src/otx/algo/classification/backbones/__init__.py
@@ -5,7 +5,6 @@
 
 from .otx_efficientnet import OTXEfficientNet
 from .otx_efficientnet_v2 import OTXEfficientNetV2
-# from .otx_mobilenet_v3 import OTXMobileNetV3
 from .mobilenet_v3 import OTXMobileNetV3
 
 __all__ = ["OTXEfficientNet", "OTXEfficientNetV2", "OTXMobileNetV3"]
diff --git a/src/otx/algo/classification/backbones/otx_mobilenet_v3.py b/src/otx/algo/classification/backbones/otx_mobilenet_v3.py
diff --git a/src/otx/algo/classification/classifier/base_classifier.py b/src/otx/algo/classification/classifier/base_classifier.py
@@ -212,8 +212,8 @@ def _forward_explain(self, images: torch.Tensor) -> dict[str, torch.Tensor | lis
         pred_results = self.head._get_predictions(logits)  # noqa: SLF001
         # H-Label Classification Case
         if isinstance(pred_results, dict):
-            scores = pred_results["pred_scores"].unbind(0)
-            preds = pred_results["pred_labels"].unbind(0)
+            scores = pred_results["pred_scores"]
+            preds = pred_results["pred_labels"]
         else:
             scores = pred_results.unbind(0)
             preds = logits.argmax(-1, keepdim=True).unbind(0)

diff --git a/src/otx/algo/classification/heads/hlabel_cls_head.py b/src/otx/algo/classification/heads/hlabel_cls_head.py
@@ -103,7 +103,7 @@ def loss(self, feats: tuple[torch.Tensor], labels: torch.Tensor, **kwargs) -> to
                 head_gt = head_gt[valid_mask]
                 if len(head_gt) > 0:
                     head_logits = head_logits[valid_mask, :]
-                    loss_score += self.multiclass_loss(head_logits, head_gt) / head_logits.size(0)
+                    loss_score += self.multiclass_loss(head_logits, head_gt)
                     num_effective_heads_in_batch += 1
 
         if num_effective_heads_in_batch > 0: