From a4abbed43d5e83553cb0231e4ab673735d12d426 Mon Sep 17 00:00:00 2001
From: Songki Choi <songki.choi@intel.com>
Date: Mon, 6 Nov 2023 18:51:18 +0900
Subject: [PATCH] Fix h-label loss normalization issue w/ exclusive label group
 of singe label (#2604)

* Fix h-label loss normalization issue w/ exclusive label group with signle label

* Fix non-linear version

---------
Signed-off-by: Songki Choi <songki.choi@intel.com>
---
 CHANGELOG.md                                                    | 1 +
 .../mmcls/models/heads/custom_hierarchical_linear_cls_head.py   | 2 +-
 .../models/heads/custom_hierarchical_non_linear_cls_head.py     | 2 +-
 3 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5f7609c4e18..5cee6309717 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -13,6 +13,7 @@ All notable changes to this project will be documented in this file.
 
 - Fix IBLoss enablement with DeiT-Tiny when class incremental training (<https://github.com/openvinotoolkit/training_extensions/pull/2595>)
 - Fix mmcls bug not wrapping model in DataParallel on CPUs (<https://github.com/openvinotoolkit/training_extensions/pull/2601>)
+- Fix h-label loss normalization issue w/ exclusive label group of singe label (<https://github.com/openvinotoolkit/training_extensions/pull/2604>)
 
 ## \[v1.4.3\]
 
diff --git a/src/otx/algorithms/classification/adapters/mmcls/models/heads/custom_hierarchical_linear_cls_head.py b/src/otx/algorithms/classification/adapters/mmcls/models/heads/custom_hierarchical_linear_cls_head.py
index 6776756bb61..3e0de200be2 100644
--- a/src/otx/algorithms/classification/adapters/mmcls/models/heads/custom_hierarchical_linear_cls_head.py
+++ b/src/otx/algorithms/classification/adapters/mmcls/models/heads/custom_hierarchical_linear_cls_head.py
@@ -105,7 +105,7 @@ def forward_train(self, cls_score, gt_label, **kwargs):
                     losses["loss"] += multiclass_loss
                     num_effective_heads_in_batch += 1
 
-        if self.hierarchical_info["num_multiclass_heads"] > 1:
+        if num_effective_heads_in_batch > 0:
             losses["loss"] /= num_effective_heads_in_batch
 
         if self.compute_multilabel_loss:
diff --git a/src/otx/algorithms/classification/adapters/mmcls/models/heads/custom_hierarchical_non_linear_cls_head.py b/src/otx/algorithms/classification/adapters/mmcls/models/heads/custom_hierarchical_non_linear_cls_head.py
index 5397818fbf3..69ea7bb1476 100644
--- a/src/otx/algorithms/classification/adapters/mmcls/models/heads/custom_hierarchical_non_linear_cls_head.py
+++ b/src/otx/algorithms/classification/adapters/mmcls/models/heads/custom_hierarchical_non_linear_cls_head.py
@@ -135,7 +135,7 @@ def forward_train(self, cls_score, gt_label, **kwargs):
                     losses["loss"] += multiclass_loss
                     num_effective_heads_in_batch += 1
 
-        if self.hierarchical_info["num_multiclass_heads"] > 1:
+        if num_effective_heads_in_batch > 0:
             losses["loss"] /= num_effective_heads_in_batch
 
         if self.compute_multilabel_loss: