✨ [New] wandb, progress class for handle proccess

requirements.txt

@@ -1,4 +1,5 @@
 einops
+graphviz
 hydra-core
 loguru
 numpy
@@ -9,4 +10,5 @@ requests
 rich
 torch
 torchvision
-tqdm
+tqdm
+wandb

yolo/tools/log_helper.py

@@ -12,32 +12,39 @@ Example:
 """
 
 import sys
-from typing import List
+from typing import Dict, List
 
+import wandb
+import wandb.errors
 from loguru import logger
 from rich.console import Console
 from rich.progress import BarColumn, Progress, TextColumn, TimeRemainingColumn
 from rich.table import Table
+from torch import Tensor
 
-from yolo.config.config import YOLOLayer
+from yolo.config.config import Config, YOLOLayer
 
 
 def custom_logger():
     logger.remove()
     logger.add(
         sys.stderr,
-        format="<green>{time:MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <level>{message}</level>",
+        format="<fg #003385>[{time:MM/DD HH:mm:ss}]</fg #003385><level>{level: ^8}</level>| <level>{message}</level>",
     )
 
 
 class CustomProgress:
-    def __init__(self):
+    def __init__(self, cfg: Config, use_wandb: bool = False):
         self.progress = Progress(
             TextColumn("[progress.description]{task.description}"),
             BarColumn(bar_width=None),
             TextColumn("{task.completed}/{task.total}"),
             TimeRemainingColumn(),
         )
+        self.use_wandb = use_wandb
+        if self.use_wandb:
+            wandb.errors.term._log = custom_wandb_log
+            self.wandb = wandb.init(project="YOLO", resume="allow", mode="online", dir="runs", name=cfg.name)
 
     def start_train(self, num_epochs: int):
         self.task_epoch = self.progress.add_task("[cyan]Epochs", total=num_epochs)
@@ -45,19 +52,34 @@ class CustomProgress:
     def one_epoch(self):
         self.progress.update(self.task_epoch, advance=1)
 
+    def finish_epoch(self):
+        self.wandb.finish()
+
     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 one_batch(self, loss_dict: Dict[str, Tensor]):
+        if self.use_wandb:
+            for loss_name, loss_value in loss_dict.items():
+                self.wandb.log({f"Loss/{loss_name}": loss_value})
+
+        loss_str = "Loss"
+        for loss_name, loss_val in loss_dict.items():
+            loss_str += f" {loss_name[:-4]}: {loss_val:.2f} |"
+
+        self.progress.update(self.batch_task, advance=1, description=f"[green]Batches [white]{loss_str}")
 
     def finish_batch(self):
         self.progress.remove_task(self.batch_task)
 
 
+def custom_wandb_log(string="", level=int, newline=True, repeat=True, prefix=True, silent=False):
+    if silent:
+        return
+    for line in string.split("\n"):
+        logger.opt(raw=not newline).info("🌐 " + line)
+
+
 def log_model(model: List[YOLOLayer]):
     console = Console()
     table = Table(title="Model Layers")

yolo/tools/trainer.py

@@ -21,6 +21,7 @@ class Trainer:
         self.optimizer = get_optimizer(model.parameters(), train_cfg.optimizer)
         self.scheduler = get_scheduler(self.optimizer, train_cfg.scheduler)
         self.loss_fn = get_loss_function(cfg)
+        self.progress = CustomProgress(cfg, use_wandb=True)
 
         if getattr(train_cfg.ema, "enabled", False):
             self.ema = EMA(model, decay=train_cfg.ema.decay)
@@ -42,21 +43,21 @@ class Trainer:
 
         return loss.item(), loss_item
 
-    def train_one_epoch(self, dataloader, progress: CustomProgress):
+    def train_one_epoch(self, dataloader):
         self.model.train()
         total_loss = 0
-        progress.start_batch(len(dataloader))
+        self.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)
+            self.progress.one_batch(loss_each)
 
         if self.scheduler:
             self.scheduler.step()
 
-        progress.finish_batch()
+        self.progress.finish_batch()
         return total_loss / len(dataloader)
 
     def save_checkpoint(self, epoch: int, filename="checkpoint.pt"):
@@ -73,14 +74,13 @@ class Trainer:
 
     def train(self, dataloader, num_epochs):
         logger.info("🚄 Start Training!")
-        progress = CustomProgress()
 
-        with progress.progress:
-            progress.start_train(num_epochs)
+        with self.progress.progress:
+            self.progress.start_train(num_epochs)
             for epoch in range(num_epochs):
 
-                epoch_loss = self.train_one_epoch(dataloader, progress)
-                progress.one_epoch()
+                epoch_loss = self.train_one_epoch(dataloader, self.progress)
+                self.progress.one_epoch()
 
                 logger.info(f"Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}")
                 if (epoch + 1) % 5 == 0:

yolo/utils/loss.py

@@ -1,5 +1,4 @@
-import time
-from typing import Any, List, Tuple
+from typing import Any, Dict, List, Tuple
 
 import torch
 import torch.nn.functional as F
@@ -169,7 +168,7 @@ class DualLoss:
         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]]:
+    def __call__(self, predicts: List[Tensor], targets: Tensor) -> Tuple[Tensor, Dict[str, Tensor]]:
         targets[:, :, 1:] = targets[:, :, 1:] * self.loss.scale_up
 
         # TODO: Need Refactor this region, make it flexible!
@@ -177,12 +176,13 @@ class DualLoss:
         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)
+        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)
+        return loss_sum, loss_dict
 
 
 def get_loss_function(cfg: Config) -> YOLOLoss: