🔨 [Update] Loss function and get_loss_func

yolo/tools/trainer.py

@@ -18,7 +18,7 @@ class Trainer:
         self.device = device
         self.optimizer = get_optimizer(model.parameters(), train_cfg.optimizer)
         self.scheduler = get_scheduler(self.optimizer, train_cfg.scheduler)
-        self.loss_fn = get_loss_function()
+        self.loss_fn = get_loss_function(cfg)
 
         if train_cfg.ema.get("enabled", False):
             self.ema = EMA(model, decay=train_cfg.ema.decay)

yolo/utils/loss.py

@@ -17,10 +17,6 @@ from yolo.tools.bbox_helper import (
 )
 
 
-def get_loss_function(*args, **kwargs):
-    raise NotImplementedError
-
-
 class BCELoss(nn.Module):
     def __init__(self) -> None:
         super().__init__()
@@ -162,5 +158,11 @@ class YOLOLoss:
         ## -- DFL -- ##
         loss_dfl = self.dfl(predicts_anc, targets_bbox, valid_masks, box_norm, cls_norm)
 
-        logger.info("Loss IoU: {:.5f}, DFL: {:.5f}, CLS: {:.5f}", loss_iou, loss_dfl, loss_cls)
-        return loss_iou, loss_dfl, loss_cls
+        loss_sum = loss_iou * 0.5 + loss_dfl * 1.5 + loss_cls * 0.5
+        return loss_sum, (loss_iou, loss_dfl, loss_cls)
+
+
+def get_loss_function(cfg: Config) -> YOLOLoss:
+    loss_function = YOLOLoss(cfg)
+    logger.info("✅ Success load loss function")
+    return loss_function