🐛 [Fix] Loss scale, scale back with batch_size

yolo/tools/loss_functions.py

@@ -124,12 +124,11 @@ class DualLoss:
         aux_iou, aux_dfl, aux_cls = self.loss(aux_predicts, targets)
         main_iou, main_dfl, main_cls = self.loss(main_predicts, targets)
 
-        loss_dict = {
-            "BoxLoss": self.iou_rate * (aux_iou * self.aux_rate + main_iou),
-            "DFLoss": self.dfl_rate * (aux_dfl * self.aux_rate + main_dfl),
-            "BCELoss": self.cls_rate * (aux_cls * self.aux_rate + main_cls),
-        }
-        loss_sum = sum(list(loss_dict.values())) / len(loss_dict)
+        BoxLoss = self.iou_rate * (aux_iou * self.aux_rate + main_iou)
+        DFLoss = self.dfl_rate * (aux_dfl * self.aux_rate + main_dfl)
+        BCELoss = self.cls_rate * (aux_cls * self.aux_rate + main_cls)
+        loss_sum = (BoxLoss + DFLoss + BCELoss) / 3
+        loss_dict = dict(BoxLoss=BoxLoss.detach(), DFLoss=DFLoss.detach(), BCELoss=BCELoss.detach())
         return loss_sum, loss_dict
 
 

yolo/tools/solver.py

@@ -66,7 +66,7 @@ class ModelTrainer:
             self.ema = None
         self.scaler = GradScaler()
 
-    def train_one_batch(self, images: Tensor, targets: Tensor):
+    def train_one_batch(self, images: Tensor, targets: Tensor, batch_size: int):
         images, targets = images.to(self.device), targets.to(self.device)
         self.optimizer.zero_grad()
 
@@ -75,7 +75,7 @@ class ModelTrainer:
             aux_predicts = self.vec2box(predicts["AUX"])
             main_predicts = self.vec2box(predicts["Main"])
             loss, loss_item = self.loss_fn(aux_predicts, main_predicts, targets)
-
+        loss *= batch_size
         self.scaler.scale(loss).backward()
         self.scaler.unscale_(self.optimizer)
         torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=10.0)
@@ -91,7 +91,7 @@ class ModelTrainer:
         self.optimizer.next_epoch(len(dataloader))
         for batch_size, images, targets, *_ in dataloader:
             self.optimizer.next_batch()
-            loss_each = self.train_one_batch(images, targets)
+            loss_each = self.train_one_batch(images, targets, batch_size)
 
             for loss_name, loss_val in loss_each.items():
                 if self.use_ddp:  # collecting loss for each batch