[!HOTFIX] Divide bounding box predictions by scaler values

CHANGELOG.md

@@ -4,7 +4,7 @@
 
 - Add instance number logging feature for detection task by `@hglee98` in [PR 577](https://github.com/Nota-NetsPresso/netspresso-trainer/pull/577)
 - Add Precision and Recall metric for detection task by `@hglee98` in [PR 579](https://github.com/Nota-NetsPresso/netspresso-trainer/pull/579)
-- Add YOLOv9 by `@hglee98` in [PR 585](https://github.com/Nota-NetsPresso/netspresso-trainer/pull/585), [PR 586](https://github.com/Nota-NetsPresso/netspresso-trainer/pull/586), [PR 592](https://github.com/Nota-NetsPresso/netspresso-trainer/pull/592), [PR 593](https://github.com/Nota-NetsPresso/netspresso-trainer/pull/593), [PR 595](https://github.com/Nota-NetsPresso/netspresso-trainer/pull/595), [PR 589](https://github.com/Nota-NetsPresso/netspresso-trainer/pull/589), [PR 590](https://github.com/Nota-NetsPresso/netspresso-trainer/pull/590), [PR 597](https://github.com/Nota-NetsPresso/netspresso-trainer/pull/597), [PR 598](https://github.com/Nota-NetsPresso/netspresso-trainer/pull/598), [PR 601](https://github.com/Nota-NetsPresso/netspresso-trainer/pull/601)
+- Add YOLOv9 by `@hglee98` in [PR 585](https://github.com/Nota-NetsPresso/netspresso-trainer/pull/585), [PR 586](https://github.com/Nota-NetsPresso/netspresso-trainer/pull/586), [PR 592](https://github.com/Nota-NetsPresso/netspresso-trainer/pull/592), [PR 593](https://github.com/Nota-NetsPresso/netspresso-trainer/pull/593), [PR 595](https://github.com/Nota-NetsPresso/netspresso-trainer/pull/595), [PR 589](https://github.com/Nota-NetsPresso/netspresso-trainer/pull/589), [PR 590](https://github.com/Nota-NetsPresso/netspresso-trainer/pull/590), [PR 597](https://github.com/Nota-NetsPresso/netspresso-trainer/pull/597), [PR 598](https://github.com/Nota-NetsPresso/netspresso-trainer/pull/598), [PR 601](https://github.com/Nota-NetsPresso/netspresso-trainer/pull/601), [PR 602](https://github.com/Nota-NetsPresso/netspresso-trainer/pull/602)
 
 ## Bug Fixes:
 

src/netspresso_trainer/losses/detection/yolov9.py

@@ -357,12 +357,8 @@ def get_output(self, output, anchor_grid, scaler):
         return pred_class_logits, pred_bbox_anchor, pred_bbox_reg
 
     def forward(self, out: Union[List, Dict], target: Dict) -> Tensor:
-        if isinstance(out['pred'], Dict):
-            aux_out = out['pred']['aux_outputs']
-            out = out['pred']['outputs']
-        else:
-            out = out['pred']
-            aux_out = None
+        # aux_out = out['pred']['aux_outputs'] if isinstance(out['pred'], Dict) else None
+        out = out['pred']['outputs'] if isinstance(out['pred'], Dict) else out['pred']
         device = out[0][0].device
         self.num_classes = target['num_classes']
         img_size = target['img_size']
@@ -381,8 +377,10 @@ def forward(self, out: Union[List, Dict], target: Dict) -> Tensor:
 
         dfl = DFLoss(anchor_grid, scaler, self.reg_max)
         preds_cls, preds_anc, preds_box = self.get_output(out, anchor_grid=anchor_grid, scaler=scaler)
-        if aux_out:
-            aux_preds_cls, aux_preds_anc, aux_preds_box = self.get_output(aux_out, anchor_grid=anchor_grid, scaler=scaler)
+        # if aux_out:
+        #     aux_preds_cls, aux_preds_anc, aux_preds_box = self.get_output(aux_out, anchor_grid=anchor_grid, scaler=scaler)
+        # else:
+        #     aux_preds_cls, aux_preds_anc, aux_preds_box = None, None, None
 
         matcher = BoxMatcher(self.num_classes, anchor_grid)
         max_samples = max([len(t['labels']) for t in target])
@@ -397,19 +395,21 @@ def forward(self, out: Union[List, Dict], target: Dict) -> Tensor:
 
         align_targets, valid_masks = matcher(labels, (preds_cls.detach(), preds_box.detach()))
         targets_cls, targets_bbox = self.separate_anchor(align_targets, scaler)
+        preds_box = preds_box / scaler[None, :, None]
         cls_norm = max(targets_cls.sum(), 1)
         box_norm = targets_cls.sum(-1)[valid_masks]
-        if aux_out:
-            aux_align_targets, aux_valid_masks = matcher(labels, (aux_preds_cls.detach(), aux_preds_box.detach()))
-            aux_targets_cls, aux_targets_bbox = self.separate_anchor(aux_align_targets, scaler)
-            aux_cls_norm = max(aux_targets_cls.sum(), 1)
-            aux_box_norm = aux_targets_cls.sum(-1)[aux_valid_masks]
-            ## -- CLS -- ##
-            aux_loss_cls = self.cls(aux_preds_cls, aux_targets_cls, aux_cls_norm)
-            ## -- IOU -- ##
-            aux_loss_iou = self.iou(aux_preds_box, aux_targets_bbox, aux_valid_masks, aux_box_norm, aux_cls_norm)
-            ## -- DFL -- ##
-            aux_loss_dfl = dfl(aux_preds_anc, aux_targets_bbox, aux_valid_masks, aux_box_norm, aux_cls_norm)
+        # if aux_out:
+        #     aux_align_targets, aux_valid_masks = matcher(labels, (aux_preds_cls.detach(), aux_preds_box.detach()))
+        #     aux_preds_box = aux_preds_box / scaler[None, :, None]
+        #     aux_targets_cls, aux_targets_bbox = self.separate_anchor(aux_align_targets, scaler)
+        #     aux_cls_norm = max(aux_targets_cls.sum(), 1)
+        #     aux_box_norm = aux_targets_cls.sum(-1)[aux_valid_masks]
+        #     ## -- CLS -- ##
+        #     aux_loss_cls = self.cls(aux_preds_cls, aux_targets_cls, aux_cls_norm)
+        #     ## -- IOU -- ##
+        #     aux_loss_iou = self.iou(aux_preds_box, aux_targets_bbox, aux_valid_masks, aux_box_norm, aux_cls_norm)
+        #     ## -- DFL -- ##
+        #     aux_loss_dfl = dfl(aux_preds_anc, aux_targets_bbox, aux_valid_masks, aux_box_norm, aux_cls_norm)
 
 
         ## -- CLS -- ##
@@ -420,10 +420,10 @@ def forward(self, out: Union[List, Dict], target: Dict) -> Tensor:
         loss_dfl = dfl(preds_anc, targets_bbox, valid_masks, box_norm, cls_norm)
 
         # TODO: loss weights should be controlled by config
-        if aux_out:
-            total_loss = 0.5 * (self.aux_rate * aux_loss_cls + loss_cls) + 1.5 * (self.aux_rate * aux_loss_dfl + loss_dfl) + 7.5 * (self.aux_rate * aux_loss_iou + loss_iou)
-        else:
-            total_loss = 0.5 * loss_cls + 1.5 * loss_dfl + 7.5 * loss_iou
+        # if aux_out:
+        #     total_loss = 0.5 * (self.aux_rate * aux_loss_cls + loss_cls) + 1.5 * (self.aux_rate * aux_loss_dfl + loss_dfl) + 7.5 * (self.aux_rate * aux_loss_iou + loss_iou)
+        # else:
+        total_loss = 0.5 * loss_cls + 1.5 * loss_dfl + 7.5 * loss_iou
         return total_loss
 
     def separate_anchor(self, anchors, scaler):

src/netspresso_trainer/models/heads/detection/experimental/yolo_head.py

@@ -109,15 +109,15 @@ def __init__(self,
             params.use_group,
         )
 
-        if params.use_aux_loss:
-            self.aux_heads = self._build_heads(
-                intermediate_features_dim,
-                params.act_type,
-                params.reg_max,
-                params.use_group,
-            )
-        else:
-            self.aux_heads = None
+        # if params.use_aux_loss:
+        #     self.aux_heads = self._build_heads(
+        #         intermediate_features_dim,
+        #         params.act_type,
+        #         params.reg_max,
+        #         params.use_group,
+        #     )
+        # else:
+        #     self.aux_heads = None
 
     def _validate_params(self, params: DictConfig) -> None:
         required_params = ['act_type', 'use_group', 'reg_max', 'num_anchors', 'use_aux_loss']
@@ -153,9 +153,9 @@ def forward(self, x_in: Union[List[Tensor], Dict], targets: Optional[Tensor] = N
         else:
             aux_in = None
         outputs = [head(x) for head, x in zip(self.heads, x_in)]
-        if self.training and self.aux_heads:
-            aux_outputs = [head(x) for head, x in zip(self.aux_heads, aux_in)]
-            outputs = {"outputs": outputs, "aux_outputs": aux_outputs}
+        # if self.training and self.aux_heads:
+        #     aux_outputs = [head(x) for head, x in zip(self.aux_heads, aux_in)]
+        #     outputs = {"outputs": outputs, "aux_outputs": aux_outputs}
         return ModelOutput(pred=outputs)
 
 def yolo_detection_head(num_classes, intermediate_features_dim, conf_model_head, **kwargs):