🔀 [Merge] branch 'DATASET'

LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2024 Kin-Yiu, Wong
+Copyright (c) 2024 Kin-Yiu, Wong and Hao-Tang, Tsui
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

yolo/config/config.py

@@ -72,7 +72,7 @@ class MatcherConfig:
 @dataclass
 class LossConfig:
     objective: List[List]
-    aux: bool
+    aux: Union[bool, float]
     matcher: MatcherConfig
 
 

yolo/config/hyper/default.yaml

@@ -1,7 +1,7 @@
 data:
-  batch_size: 8 
+  batch_size: 16 
   shuffle: True
-  num_workers: 4
+  num_workers: 16
   pin_memory: True
   class_num: 80
   image_size: [640, 640]
@@ -13,11 +13,11 @@ train:
       weight_decay: 0.0001
   loss:
     objective:
-      - [BCELoss, 0.1] 
-      - [BoxLoss, 0.1]  
-      - [DFLoss, 0.1]
+      BCELoss: 0.5
+      BoxLoss: 7.5  
+      DFLoss: 1.5
     aux:
-      True
+      0.25
     matcher:
       iou: CIoU
       topk: 10

yolo/tools/log_helper.py

@@ -16,6 +16,7 @@ from typing import List
 
 from loguru import logger
 from rich.console import Console
+from rich.progress import BarColumn, Progress, TextColumn, TimeRemainingColumn
 from rich.table import Table
 
 from yolo.config.config import YOLOLayer
@@ -29,6 +30,34 @@ def custom_logger():
     )
 
 
+class CustomProgress:
+    def __init__(self):
+        self.progress = Progress(
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(bar_width=None),
+            TextColumn("{task.completed}/{task.total}"),
+            TimeRemainingColumn(),
+        )
+
+    def start_train(self, num_epochs: int):
+        self.task_epoch = self.progress.add_task("[cyan]Epochs", total=num_epochs)
+
+    def one_epoch(self):
+        self.progress.update(self.task_epoch, advance=1)
+
+    def start_batch(self, num_batches):
+        self.batch_task = self.progress.add_task("[green]Batches", total=num_batches)
+
+    def one_batch(self, loss_each):
+        loss_iou, loss_dfl, loss_cls = loss_each
+        # TODO: make it flexible? if need add more loss
+        loss_str = f"Loss IoU: {loss_iou:.3f}, DFL: {loss_dfl:.3f}, CLS: {loss_cls:.3f}"
+        self.progress.update(self.batch_task, advance=1, description=f"[green]Batches {loss_str}")
+
+    def finish_batch(self):
+        self.progress.remove_task(self.batch_task)
+
+
 def log_model(model: List[YOLOLayer]):
     console = Console()
     table = Table(title="Model Layers")

yolo/tools/trainer.py

@@ -1,11 +1,13 @@
 import torch
 from loguru import logger
 from torch import Tensor
+
+# TODO: We may can't use CUDA?
 from torch.cuda.amp import GradScaler, autocast
-from tqdm import tqdm
 
 from yolo.config.config import Config, TrainConfig
 from yolo.model.yolo import YOLO
+from yolo.tools.log_helper import CustomProgress
 from yolo.tools.model_helper import EMA, get_optimizer, get_scheduler
 from yolo.utils.loss import get_loss_function
 
@@ -26,16 +28,13 @@ class Trainer:
             self.ema = None
         self.scaler = GradScaler()
 
-    def train_one_batch(self, data: Tensor, targets: Tensor, progress: tqdm):
+    def train_one_batch(self, data: Tensor, targets: Tensor):
         data, targets = data.to(self.device), targets.to(self.device)
         self.optimizer.zero_grad()
 
         with autocast():
             outputs = self.model(data)
             loss, loss_item = self.loss_fn(outputs, targets)
-            loss_iou, loss_dfl, loss_cls = loss_item
-
-        progress.set_description(f"Loss IoU: {loss_iou:.5f}, DFL: {loss_dfl:.5f}, CLS: {loss_cls:.5f}")
 
         self.scaler.scale(loss).backward()
         self.scaler.step(self.optimizer)
@@ -43,17 +42,21 @@ class Trainer:
 
         return loss.item(), loss_item
 
-        return loss.item()
-
-    def train_one_epoch(self, dataloader):
+    def train_one_epoch(self, dataloader, progress: CustomProgress):
         self.model.train()
         total_loss = 0
-        with tqdm(dataloader, desc="Training") as progress:
-            for data, targets in progress:
-                loss = self.train_one_batch(data, targets, progress)
-                total_loss += loss
-            if self.scheduler:
-                self.scheduler.step()
+        progress.start_batch(len(dataloader))
+
+        for data, targets in dataloader:
+            loss, loss_each = self.train_one_batch(data, targets)
+
+            total_loss += loss
+            progress.one_batch(loss_each)
+
+        if self.scheduler:
+            self.scheduler.step()
+
+        progress.finish_batch()
         return total_loss / len(dataloader)
 
     def save_checkpoint(self, epoch: int, filename="checkpoint.pt"):
@@ -69,9 +72,16 @@ class Trainer:
         torch.save(checkpoint, filename)
 
     def train(self, dataloader, num_epochs):
-        logger.info("start train")
-        for epoch in range(num_epochs):
-            epoch_loss = self.train_one_epoch(dataloader)
-            logger.info(f"Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}")
-            if (epoch + 1) % 5 == 0:
-                self.save_checkpoint(epoch, f"checkpoint_epoch_{epoch+1}.pth")
+        logger.info("🚄 Start Training!")
+        progress = CustomProgress()
+
+        with progress.progress:
+            progress.start_train(num_epochs)
+            for epoch in range(num_epochs):
+
+                epoch_loss = self.train_one_epoch(dataloader, progress)
+                progress.one_epoch()
+
+                logger.info(f"Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}")
+                if (epoch + 1) % 5 == 0:
+                    self.save_checkpoint(epoch, f"checkpoint_epoch_{epoch+1}.pth")

yolo/utils/dataloader.py

@@ -178,13 +178,15 @@ class YoloDataLoader(DataLoader):
                 - A tensor of batched images.
                 - A list of tensors, each corresponding to bboxes for each image in the batch.
         """
-        batch_size, target_size = len(batch), [item[1].size(0) for item in batch]
-        batch_targets = torch.zeros(batch_size, max(target_size), 5)
-        images = []
-        for idx, (image, target) in enumerate(batch):
-            images.append(image)
-            batch_targets[idx, : target_size[idx]] = target
-        batch_images = torch.stack(images)
+        batch_size = len(batch)
+        target_sizes = [item[1].size(0) for item in batch]
+        # TODO: Improve readability of these proccess
+        batch_targets = torch.zeros(batch_size, max(target_sizes), 5)
+        for idx, target_size in enumerate(target_sizes):
+            batch_targets[idx, :target_size] = batch[idx][1]
+
+        batch_images = torch.stack([item[0] for item in batch])
+
         return batch_images, batch_targets
 
 

yolo/utils/loss.py

@@ -15,6 +15,7 @@ from yolo.tools.bbox_helper import (
     make_anchor,
     transform_bbox,
 )
+from yolo.tools.module_helper import make_chunk
 
 
 class BCELoss(nn.Module):
@@ -135,14 +136,10 @@ class YOLOLoss:
         anchors_box = anchors_box / self.scaler[None, :, None]
         return anchors_cls, anchors_box
 
-    @torch.autocast("cuda" if torch.cuda.is_available() else "cpu")
     def __call__(self, predicts: List[Tensor], targets: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
         # Batch_Size x (Anchor + Class) x H x W
         # TODO: check datatype, why targets has a little bit error with origin version
-        predicts, predicts_anc = self.parse_predicts(predicts[0])
-        # TODO: Refactor this operator
-        # targets = self.parse_targets(targets, batch_size=predicts.size(0))
-        targets[:, :, 1:] = targets[:, :, 1:] * self.scale_up
+        predicts, predicts_anc = self.parse_predicts(predicts)
 
         align_targets, valid_masks = self.matcher(targets, predicts)
         # calculate loss between with instance and predict
@@ -160,11 +157,35 @@ class YOLOLoss:
         ## -- DFL -- ##
         loss_dfl = self.dfl(predicts_anc, targets_bbox, valid_masks, box_norm, cls_norm)
 
-        loss_sum = loss_iou * 0.5 + loss_dfl * 1.5 + loss_cls * 0.5
-        return loss_sum, (loss_iou, loss_dfl, loss_cls)
+        return loss_iou, loss_dfl, loss_cls
+
+
+class DualLoss:
+    def __init__(self, cfg: Config) -> None:
+        self.loss = YOLOLoss(cfg)
+        self.aux_rate = cfg.hyper.train.loss.aux
+
+        self.iou_rate = cfg.hyper.train.loss.objective["BoxLoss"]
+        self.dfl_rate = cfg.hyper.train.loss.objective["DFLoss"]
+        self.cls_rate = cfg.hyper.train.loss.objective["BCELoss"]
+
+    def __call__(self, predicts: List[Tensor], targets: Tensor) -> Tuple[Tensor, Tuple[Tensor]]:
+        targets[:, :, 1:] = targets[:, :, 1:] * self.loss.scale_up
+
+        # TODO: Need Refactor this region, make it flexible!
+        predicts = make_chunk(predicts[0], 2)
+        aux_iou, aux_dfl, aux_cls = self.loss(predicts[0], targets)
+        main_iou, main_dfl, main_cls = self.loss(predicts[1], targets)
+
+        loss_iou = self.iou_rate * (aux_iou * self.aux_rate + main_iou)
+        loss_dfl = self.dfl_rate * (aux_dfl * self.aux_rate + main_dfl)
+        loss_cls = self.cls_rate * (aux_cls * self.aux_rate + main_cls)
+
+        loss = (loss_iou + loss_dfl + loss_cls) / 3
+        return loss, (loss_iou, loss_dfl, loss_cls)
 
 
 def get_loss_function(cfg: Config) -> YOLOLoss:
-    loss_function = YOLOLoss(cfg)
+    loss_function = DualLoss(cfg)
     logger.info("✅ Success load loss function")
     return loss_function