✨ [Add] General config for global settings

examples/example_train.py

@@ -9,8 +9,7 @@ project_root = Path(__file__).resolve().parent.parent
 sys.path.append(str(project_root))
 
 from yolo.config.config import Config
-from yolo.model.yolo import get_model
-from yolo.tools.log_helper import custom_logger
+from yolo.tools.log_helper import custom_logger, get_valid_folder
 from yolo.tools.trainer import Trainer
 from yolo.utils.dataloader import get_dataloader
 from yolo.utils.get_dataset import prepare_dataset
@@ -18,18 +17,17 @@ from yolo.utils.get_dataset import prepare_dataset
 
 @hydra.main(config_path="../yolo/config", config_name="config", version_base=None)
 def main(cfg: Config):
+    custom_logger()
+    save_path = get_valid_folder(cfg.hyper.general, cfg.name)
     if cfg.download.auto:
         prepare_dataset(cfg.download)
 
     dataloader = get_dataloader(cfg)
-    model = get_model(cfg)
     # TODO: get_device or rank, for DDP mode
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    trainer = Trainer(model, cfg, device)
-    trainer.train(dataloader, 10)
+    trainer = Trainer(cfg, save_path, device)
+    trainer.train(dataloader, cfg.hyper.train.epoch)
 
 
 if __name__ == "__main__":
-    custom_logger()
     main()

yolo/config/config.py

@@ -25,11 +25,10 @@ class Download:
 @dataclass
 class DataLoaderConfig:
     batch_size: int
+    class_num: int
+    image_size: List[int]
     shuffle: bool
-    num_workers: int
     pin_memory: bool
-    image_size: List[int]
-    class_num: int
 
 
 @dataclass
@@ -85,8 +84,22 @@ class TrainConfig:
     loss: LossConfig
 
 
+@dataclass
+class GeneralConfig:
+    out_path: str
+    task: str
+    device: Union[str, int, List[int]]
+    cpu_num: int
+    use_wandb: bool
+    lucky_number: 10
+    exist_ok: bool
+    resume_train: bool
+    use_TensorBoard: bool
+
+
 @dataclass
 class HyperConfig:
+    general: GeneralConfig
     data: DataLoaderConfig
     train: TrainConfig
 

yolo/config/hyper/default.yaml

@@ -1,10 +1,19 @@
+general:
+  out_path: runs
+  task: train
+  deivce: [0]
+  cpu_num: 16
+  use_wandb: False
+  lucky_number: 10
+  exist_ok: True
+  resume_train: False
+  use_TensorBoard: False
 data:
   batch_size: 16 
-  shuffle: True
-  num_workers: 16
-  pin_memory: True
   class_num: 80
   image_size: [640, 640]
+  shuffle: True
+  pin_memory: True
 train:
   epoch: 10
   optimizer:

yolo/tools/log_helper.py

@@ -11,6 +11,7 @@ Example:
     custom_logger()
 """
 
+import os
 import sys
 from typing import Dict, List
 
@@ -22,19 +23,20 @@ from rich.progress import BarColumn, Progress, TextColumn, TimeRemainingColumn
 from rich.table import Table
 from torch import Tensor
 
-from yolo.config.config import Config, YOLOLayer
+from yolo.config.config import Config, GeneralConfig, YOLOLayer
 
 
 def custom_logger():
     logger.remove()
     logger.add(
         sys.stderr,
-        format="<fg #003385>[{time:MM/DD HH:mm:ss}]</fg #003385><level>{level: ^8}</level>| <level>{message}</level>",
+        colorize=True,
+        format="<fg #003385>[{time:MM/DD HH:mm:ss}]</> <level>{level: ^8}</level>| <level>{message}</level>",
     )
 
 
 class CustomProgress:
-    def __init__(self, cfg: Config, use_wandb: bool = False):
+    def __init__(self, cfg: Config, save_path: str, use_wandb: bool = False):
         self.progress = Progress(
             TextColumn("[progress.description]{task.description}"),
             BarColumn(bar_width=None),
@@ -44,18 +46,19 @@ class CustomProgress:
         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)
+            self.wandb = wandb.init(
+                project="YOLO", resume="allow", mode="online", dir=save_path, id=None, name=cfg.name
+            )
 
     def start_train(self, num_epochs: int):
         self.task_epoch = self.progress.add_task("[cyan]Epochs  [white]| Loss | Box  | DFL  | BCE  |", total=num_epochs)
 
-    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):
+    def start_one_epoch(self, num_batches, optimizer, epoch_idx):
+        if self.use_wandb:
+            lr_values = [params["lr"] for params in optimizer.param_groups]
+            lr_names = ["bias", "norm", "conv"]
+            for lr_name, lr_value in zip(lr_names, lr_values):
+                self.wandb.log({f"Learning Rate/{lr_name}": lr_value}, step=epoch_idx)
         self.batch_task = self.progress.add_task("[green]Batches", total=num_batches)
 
     def one_batch(self, loss_dict: Dict[str, Tensor]):
@@ -69,15 +72,19 @@ class CustomProgress:
 
         self.progress.update(self.batch_task, advance=1, description=f"[green]Batches [white]{loss_str}")
 
-    def finish_batch(self):
+    def finish_one_epoch(self):
         self.progress.remove_task(self.batch_task)
+        self.progress.update(self.task_epoch, advance=1)
+
+    def finish_train(self):
+        self.wandb.finish()
 
 
 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)
+        logger.opt(raw=not newline, colors=True).info("🌐 " + line)
 
 
 def log_model(model: List[YOLOLayer]):
@@ -99,3 +106,25 @@ def log_model(model: List[YOLOLayer]):
             channels = "-"
         table.add_row(str(idx), layer.layer_type, layer.tags, f"{layer_param:,}", channels)
     console.print(table)
+
+
+def get_valid_folder(general_cfg: GeneralConfig, exp_name):
+    base_path = os.path.join(general_cfg.out_path, general_cfg.task)
+    save_path = os.path.join(base_path, exp_name)
+
+    if not general_cfg.exist_ok:
+        index = 1
+        old_exp_name = exp_name
+        while os.path.isdir(save_path):
+            exp_name = f"{old_exp_name}{index}"
+            save_path = os.path.join(base_path, exp_name)
+            index += 1
+        if index > 1:
+            logger.opt(colors=True).warning(
+                f"🔀 Experiment directory exists! Changed <red>{old_exp_name}</> to <green>{exp_name}</>"
+            )
+
+    os.makedirs(save_path, exist_ok=True)
+    logger.opt(colors=True).info(f"📄 Created log folder: <u><fg #808080>{save_path}</></>")
+    logger.add(os.path.join(save_path, "output.log"), backtrace=True, diagnose=True)
+    return save_path

yolo/tools/trainer.py

@@ -6,22 +6,23 @@ from torch import Tensor
 from torch.cuda.amp import GradScaler, autocast
 
 from yolo.config.config import Config, TrainConfig
-from yolo.model.yolo import YOLO
+from yolo.model.yolo import get_model
 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
 
 
 class Trainer:
-    def __init__(self, model: YOLO, cfg: Config, device):
+    def __init__(self, cfg: Config, save_path: str, device):
         train_cfg: TrainConfig = cfg.hyper.train
+        model = get_model(cfg)
 
         self.model = model.to(device)
         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(cfg)
-        self.progress = CustomProgress(cfg, use_wandb=True)
+        self.progress = CustomProgress(cfg, save_path, use_wandb=True)
 
         if getattr(train_cfg.ema, "enabled", False):
             self.ema = EMA(model, decay=train_cfg.ema.decay)

yolo/utils/dataloader.py

@@ -160,7 +160,7 @@ class YoloDataLoader(DataLoader):
             dataset,
             batch_size=hyper.batch_size,
             shuffle=hyper.shuffle,
-            num_workers=hyper.num_workers,
+            num_workers=config.hyper.general.cpu_num,
             pin_memory=hyper.pin_memory,
             collate_fn=self.collate_fn,
         )