🔀 [Merge] branch 'TRAIN' into TEST

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
@@ -54,6 +53,7 @@ class SchedulerArgs:
 class SchedulerConfig:
     type: str
     args: SchedulerArgs
+    warmup: Dict[str, Union[str, int, float]]
 
 
 @dataclass
@@ -85,8 +85,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,17 +1,27 @@
+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
+  epoch: 500
   optimizer:
-    type: Adam
+    type: SGD
     args:
-      lr: 0.001
-      weight_decay: 0.0001
+      lr: 0.01
+      weight_decay: 0.0005
+      momentum: 0.937
   loss:
     objective:
       BCELoss: 0.5
@@ -26,10 +36,13 @@ train:
         iou: 6.0
         cls: 0.5
   scheduler:
-    type: StepLR
+    type: LinearLR
+    warmup:
+      epochs: 3.0
     args:
-      step_size: 10
-      gamma: 0.1
+      total_iters: ${hyper.train.epoch}
+      start_factor: 1
+      end_factor: 0.01
   ema:
       enabled: true
       decay: 0.995

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/model_helper.py

@@ -2,9 +2,10 @@ from typing import Any, Dict, Type
 
 import torch
 from torch.optim import Optimizer
-from torch.optim.lr_scheduler import _LRScheduler
+from torch.optim.lr_scheduler import LambdaLR, SequentialLR, _LRScheduler
 
 from yolo.config.config import OptimizerConfig, SchedulerConfig
+from yolo.model.yolo import YOLO
 
 
 class EMA:
@@ -31,21 +32,38 @@ class EMA:
             self.shadow[name].copy_(param.data)
 
 
-def get_optimizer(model_parameters, optim_cfg: OptimizerConfig) -> Optimizer:
+def get_optimizer(model: YOLO, optim_cfg: OptimizerConfig) -> Optimizer:
     """Create an optimizer for the given model parameters based on the configuration.
 
     Returns:
         An instance of the optimizer configured according to the provided settings.
     """
     optimizer_class: Type[Optimizer] = getattr(torch.optim, optim_cfg.type)
+
+    bias_params = [p for name, p in model.named_parameters() if "bias" in name]
+    norm_params = [p for name, p in model.named_parameters() if "weight" in name and "bn" in name]
+    conv_params = [p for name, p in model.named_parameters() if "weight" in name and "bn" not in name]
+
+    model_parameters = [
+        {"params": bias_params, "nestrov": True, "momentum": 0.937},
+        {"params": conv_params, "weight_decay": 0.0},
+        {"params": norm_params, "weight_decay": 1e-5},
+    ]
     return optimizer_class(model_parameters, **optim_cfg.args)
 
 
-def get_scheduler(optimizer: Optimizer, schedul_cfg: SchedulerConfig) -> _LRScheduler:
+def get_scheduler(optimizer: Optimizer, schedule_cfg: SchedulerConfig) -> _LRScheduler:
     """Create a learning rate scheduler for the given optimizer based on the configuration.
 
     Returns:
         An instance of the scheduler configured according to the provided settings.
     """
-    scheduler_class: Type[_LRScheduler] = getattr(torch.optim.lr_scheduler, schedul_cfg.type)
-    return scheduler_class(optimizer, **schedul_cfg.args)
+    scheduler_class: Type[_LRScheduler] = getattr(torch.optim.lr_scheduler, schedule_cfg.type)
+    schedule = scheduler_class(optimizer, **schedule_cfg.args)
+    if hasattr(schedule_cfg, "warmup"):
+        wepoch = schedule_cfg.warmup.epochs
+        lambda1 = lambda epoch: 0.1 + 0.9 * (epoch + 1 / wepoch) if epoch < wepoch else 1
+        lambda2 = lambda epoch: 10 - 9 * (epoch + 1 / wepoch) if epoch < wepoch else 1
+        warmup_schedule = LambdaLR(optimizer, lr_lambda=[lambda1, lambda2, lambda1])
+        schedule = SequentialLR(optimizer, schedulers=[warmup_schedule, schedule], milestones=[2])
+    return schedule

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.optimizer = get_optimizer(model, 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)
@@ -46,7 +47,6 @@ class Trainer:
     def train_one_epoch(self, dataloader):
         self.model.train()
         total_loss = 0
-        self.progress.start_batch(len(dataloader))
 
         for data, targets in dataloader:
             loss, loss_each = self.train_one_batch(data, targets)
@@ -57,7 +57,6 @@ class Trainer:
         if self.scheduler:
             self.scheduler.step()
 
-        self.progress.finish_batch()
         return total_loss / len(dataloader)
 
     def save_checkpoint(self, epoch: int, filename="checkpoint.pt"):
@@ -79,8 +78,9 @@ class Trainer:
             self.progress.start_train(num_epochs)
             for epoch in range(num_epochs):
 
+                self.progress.start_one_epoch(len(dataloader), self.optimizer, epoch)
                 epoch_loss = self.train_one_epoch(dataloader)
-                self.progress.one_epoch()
+                self.progress.finish_one_epoch()
 
                 logger.info(f"Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}")
                 if (epoch + 1) % 5 == 0:

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,
         )