From 78866031a0df6b40bc09ff24515214a0945768d3 Mon Sep 17 00:00:00 2001 From: rchuzh Date: Sat, 28 Sep 2024 13:25:30 +0800 Subject: [PATCH] fix(mmdet/base_dense_head): :bug: correct slicing issue in score softmax for single-class models The softmax function previously sliced out the background class, which caused an IndexError in single-class model export. This update removes the slicing of background class scores, fixing the issue when exporting RetinaNet with single-class configurations. - Updated `base_dense_head__predict_by_feat` to prevent IndexError when handling single-class outputs. - This change is particularly relevant for RetinaNet model exports. --- mmdeploy/codebase/mmdet/models/dense_heads/base_dense_head.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mmdeploy/codebase/mmdet/models/dense_heads/base_dense_head.py b/mmdeploy/codebase/mmdet/models/dense_heads/base_dense_head.py index 670cb10875..0cb2aea5dd 100644 --- a/mmdeploy/codebase/mmdet/models/dense_heads/base_dense_head.py +++ b/mmdeploy/codebase/mmdet/models/dense_heads/base_dense_head.py @@ -113,7 +113,7 @@ def base_dense_head__predict_by_feat( if self.use_sigmoid_cls: scores = scores.sigmoid() else: - scores = scores.softmax(-1)[:, :, :-1] + scores = scores.softmax(-1) if with_score_factors: score_factors = score_factors.permute(0, 2, 3, 1).reshape(batch_size,