✨ [Init] Trainer for training whole model!

config/config.py

@@ -14,7 +14,59 @@ class Download:
     path: str
 
 
+@dataclass
+class DataLoaderConfig:
+    batch_size: int
+    shuffle: bool
+    num_workers: int
+    pin_memory: bool
+
+
+@dataclass
+class OptimizerArgs:
+    lr: float
+    weight_decay: float
+
+
+@dataclass
+class OptimizerConfig:
+    type: str
+    args: OptimizerArgs
+
+
+@dataclass
+class SchedulerArgs:
+    step_size: int
+    gamma: float
+
+
+@dataclass
+class SchedulerConfig:
+    type: str
+    args: SchedulerArgs
+
+
+@dataclass
+class EMAConfig:
+    enabled: bool
+    decay: float
+
+
+@dataclass
+class TrainConfig:
+    optimizer: OptimizerConfig
+    scheduler: SchedulerConfig
+    ema: EMAConfig
+
+
+@dataclass
+class HyperConfig:
+    data: DataLoaderConfig
+    train: TrainConfig
+
+
 @dataclass
 class Config:
     model: Model
     download: Download
+    hyper: HyperConfig

config/hyper/default.yaml

@@ -3,3 +3,17 @@ data:
   shuffle: True
   num_workers: 4
   pin_memory: True
+train:
+  optimizer:
+    type: Adam
+    args:
+      lr: 0.001
+      weight_decay: 0.0001
+  scheduler:
+    type: StepLR
+    args:
+      step_size: 10
+      gamma: 0.1
+  ema:
+      enabled: true
+      decay: 0.995

tools/model_helper.py

@@ -0,0 +1,51 @@
+from typing import Any, Dict, Type
+
+import torch
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import _LRScheduler
+
+from config.config import OptimizerConfig, SchedulerConfig
+
+
+class EMA:
+    def __init__(self, model: torch.nn.Module, decay: float):
+        self.model = model
+        self.decay = decay
+        self.shadow = {name: param.clone().detach() for name, param in model.named_parameters()}
+
+    def update(self):
+        """Update the shadow parameters using the current model parameters."""
+        for name, param in self.model.named_parameters():
+            assert name in self.shadow, "All model parameters should have a corresponding shadow parameter."
+            new_average = (1.0 - self.decay) * param.data + self.decay * self.shadow[name]
+            self.shadow[name] = new_average.clone()
+
+    def apply_shadow(self):
+        """Apply the shadow parameters to the model."""
+        for name, param in self.model.named_parameters():
+            param.data.copy_(self.shadow[name])
+
+    def restore(self):
+        """Restore the original parameters from the shadow."""
+        for name, param in self.model.named_parameters():
+            self.shadow[name].copy_(param.data)
+
+
+def get_optimizer(model_parameters, 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)
+    return optimizer_class(model_parameters, **optim_cfg.args)
+
+
+def get_scheduler(optimizer: Optimizer, schedul_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)

tools/trainer.py

@@ -0,0 +1,63 @@
+import torch
+from loguru import logger
+from tqdm import tqdm
+
+from config.config import TrainConfig
+from model.yolo import YOLO
+from tools.model_helper import EMA, get_optimizer, get_scheduler
+from utils.loss import get_loss_function
+
+
+class Trainer:
+    def __init__(self, model: YOLO, train_cfg: TrainConfig, device):
+        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()
+
+        if train_cfg.ema.get("enabled", False):
+            self.ema = EMA(model, decay=train_cfg.ema.decay)
+        else:
+            self.ema = None
+
+    def train_one_batch(self, data, targets):
+        data, targets = data.to(self.device), targets.to(self.device)
+        self.optimizer.zero_grad()
+        outputs = self.model(data)
+        loss = self.loss_fn(outputs, targets)
+        loss.backward()
+        self.optimizer.step()
+        if self.ema:
+            self.ema.update()
+        return loss.item()
+
+    def train_one_epoch(self, dataloader):
+        self.model.train()
+        total_loss = 0
+        for data, targets in tqdm(dataloader, desc="Training"):
+            loss = self.train_one_batch(data, targets)
+            total_loss += loss
+        if self.scheduler:
+            self.scheduler.step()
+        return total_loss / len(dataloader)
+
+    def save_checkpoint(self, epoch, filename="checkpoint.pt"):
+        checkpoint = {
+            "epoch": epoch,
+            "model_state_dict": self.model.state_dict(),
+            "optimizer_state_dict": self.optimizer.state_dict(),
+        }
+        if self.ema:
+            self.ema.apply_shadow()
+            checkpoint["model_state_dict_ema"] = self.model.state_dict()
+            self.ema.restore()
+        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")

train.py

@@ -1,20 +1,27 @@
 import hydra
+import torch
 from loguru import logger
 
 from config.config import Config
 from model.yolo import get_model
 from tools.log_helper import custom_logger
-from utils.dataloader import YoloDataset
+from tools.trainer import Trainer
+from utils.dataloader import get_dataloader
 from utils.get_dataset import prepare_dataset
 
 
 @hydra.main(config_path="config", config_name="config", version_base=None)
 def main(cfg: Config):
-    dataset = YoloDataset(cfg)
     if cfg.download.auto:
         prepare_dataset(cfg.download)
 
+    dataloader = get_dataloader(cfg)
     model = get_model(cfg.model)
+    # TODO: get_device or rank, for DDP mode
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    trainer = Trainer(model, cfg.hyper.train, device)
+    trainer.train(dataloader, 10)
 
 
 if __name__ == "__main__":

utils/loss.py

@@ -0,0 +1,2 @@
+def get_loss_function(*args, **kwargs):
+    raise NotImplementedError