🔀 [Merge] branch 'main' into DATASET

.github/workflows/main.yaml

@@ -1,4 +1,4 @@
-name: YOLOv9 - Model test
+name: YOLOv9 - Model test and Code Style Check
 
 on:
   push:
@@ -8,7 +8,6 @@ on:
 
 jobs:
   build:
-
     runs-on: ubuntu-latest
 
     steps:
@@ -17,10 +16,25 @@ jobs:
       uses: actions/setup-python@v2
       with:
         python-version: 3.8
+
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
         pip install -r requirements.txt
+
+    - name: Install pre-commit
+      run: pip install pre-commit
+
+    - name: Cache pre-commit environment
+      uses: actions/cache@v2
+      with:
+        path: ~/.cache/pre-commit
+        key: ${{ runner.os }}-precommit-${{ hashFiles('**/.pre-commit-config.yaml') }}
+        restore-keys: |
+          ${{ runner.os }}-precommit-
+
+    - name: Run pre-commit (black and isort)
+      run: pre-commit run --all-files
+
     - name: Test with pytest
-      run: |
-        pytest
+      run: pytest

config/config.py

@@ -14,6 +14,57 @@ 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 Dataset:
     file_name: str
@@ -37,3 +88,4 @@ class Download:
 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

tests/test_utils/test_dataaugment.py

@@ -6,23 +6,22 @@ from PIL import Image
 from torchvision.transforms import functional as TF
 
 sys.path.append("./")
-from utils.data_augment import Compose, Mosaic, RandomHorizontalFlip
+from utils.data_augment import Compose, HorizontalFlip, Mosaic, VerticalFlip
 
 
-def test_random_horizontal_flip():
+def test_horizontal_flip():
     # Create a mock image and bounding boxes
     img = Image.new("RGB", (100, 100), color="red")
-    boxes = torch.tensor([[1, 0.1, 0.1, 0.9, 0.9]])  # class, xmin, ymin, xmax, ymax
+    boxes = torch.tensor([[1, 0.05, 0.1, 0.7, 0.9]])  # class, xmin, ymin, xmax, ymax
 
-    flip_transform = RandomHorizontalFlip(prob=1)  # Set probability to 1 to ensure flip
+    flip_transform = HorizontalFlip(prob=1)  # Set probability to 1 to ensure flip
     flipped_img, flipped_boxes = flip_transform(img, boxes)
 
     # Assert image is flipped by comparing it to a manually flipped image
     assert TF.hflip(img) == flipped_img
 
     # Assert bounding boxes are flipped correctly
-    expected_boxes = torch.tensor([[1, 0.1, 0.1, 0.9, 0.9]])
-    expected_boxes[:, [1, 3]] = 1 - expected_boxes[:, [3, 1]]
+    expected_boxes = torch.tensor([[1, 0.3, 0.1, 0.95, 0.9]])
     assert torch.allclose(flipped_boxes, expected_boxes), "Bounding boxes were not flipped correctly"
 
 
@@ -60,5 +59,5 @@ def test_mosaic():
     # Checks here would depend on the exact expected behavior of the mosaic function,
     # such as dimensions and content of the output image and boxes.
 
-    assert mosaic_img.size == (200, 200), "Mosaic image size should be doubled"
+    assert mosaic_img.size == (100, 100), "Mosaic image size should be same"
     assert len(mosaic_boxes) > 0, "Should have some bounding boxes"

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