🔀 [Merge] branch 'main' into DATASET

README.md

@@ -1,5 +1,4 @@
-# YOLOv9-MIT
-An MIT license rewrite of YOLOv9
+# YOLO: Official Implementation of YOLOv{7, 9}
 
 ![WIP](https://img.shields.io/badge/status-WIP-orange)
 > [!IMPORTANT]
@@ -7,44 +6,75 @@ An MIT license rewrite of YOLOv9
 > 
 > Use of this code is at your own risk and discretion. It is advisable to consult with the project owner before deploying or integrating into any critical systems.
 
-## Contributing
+Welcome to the official implementation of the YOLOv7 and YOLOv9. This repository will contains the complete codebase, pre-trained models, and detailed instructions for training and deploying YOLOv9.
 
-While the project's structure is still being finalized, we ask that potential contributors wait for these foundational decisions to be made. We greatly appreciate your patience and are excited to welcome contributions from the community once we are ready. Alternatively, you are welcome to propose functions that should be implemented based on the original YOLO version or suggest other enhancements! 
+## TL;DR
+- Official YOLOv9 model implementation.
+- Features real-time detection with state-of-the-art accuracy.
+<!-- - Includes pre-trained models and training scripts. -->
+- Quick train: `python examples/example_train.py`
 
-If you are interested in contributing, please keep an eye on project updates or contact us directly at [[email protected]](mailto:[email protected]) for more information.
+## Introduction
+- [**YOLOv9**: Learning What You Want to Learn Using Programmable Gradient Information](https://arxiv.org/abs/2402.13616)
+- [**YOLOv7**: Trainable bag-of-freebies sets new state-of-the-art for real-time object detectors](https://arxiv.org/abs/2207.02696)
 
-## Star History
+## Installation
+To get started with YOLOv9, clone this repository and install the required dependencies:
+```shell
+git clone [email protected]:WongKinYiu/yolov9mit.git
+cd yolov9mit
+pip install -r requirements.txt
+```
+
+<!-- 
+```
+pip install git+https://github.com/WongKinYiu/yolov9mit.git
+``` 
+-->
+
+<!-- ### Quick Start
+Run YOLOv9 on a pre-trained model with:
+
+```shell
+python examples/example_train.py hyper.data.batch_size=8
+``` -->
+
+<!-- ## Model Zoo[WIP]
+Find pre-trained models with benchmarks on various datasets in the [Model Zoo](docs/MODELS). -->
+
+## Training
+For training YOLOv9 on your dataset:
+
+Modify the configuration file data/config.yaml to point to your dataset.
+Run the training script:
 
+```shell
+python examples/example_train.py hyper.data.batch_size=8 model=v9-c
+```
+
+More customization details, or ways to modify the model can be found [HOWTO](docs/HOWTO).
+
+## Evaluation [WIP]
+Evaluate the model performance using:
+
+```shell
+python examples/examples_evaluate.py weights=v9-c.pt
+```
+
+## Contributing
+Contributions to the YOLOv9 project are welcome! See [CONTRIBUTING](docs/CONTRIBUTING.md) for how to help out.
+
+## Star History
 [![Star History Chart](https://api.star-history.com/svg?repos=WongKinYiu/yolov9mit&type=Date)](https://star-history.com/#WongKinYiu/yolov9mit&Date)
 
-## To-Do Lists
-- [ ] Project Setup
-    - [X] requirements
-    - [x] LICENSE
-    - [ ] README
-    - [x] pytests
-    - [ ] setup.py/pip install
-    - [x] log format
-    - [ ] hugging face
-- [ ] Data proccess
-    - [ ] Dataset
-        - [x] Download script
-        - [ ] Auto Download
-        - [ ] xywh, xxyy, xcyc
-    - [x] Dataloder
-        - [x] Data arugment
-- [ ] Model
-    - [ ] load model
-        - [ ] from yaml
-        - [ ] from github
-    - [x] trainer
-        - [x] train_one_iter
-        - [x] train_one_epoch
-    - [ ] DDP
-    - [x] EMA, OTA
-- [ ] Loss
-- [ ] Run
-    - [ ] train
-    - [ ] test
-    - [ ] demo
-- [x] Configuration
+## Citations
+```
+@misc{wang2024yolov9,
+      title={YOLOv9: Learning What You Want to Learn Using Programmable Gradient Information}, 
+      author={Chien-Yao Wang and I-Hau Yeh and Hong-Yuan Mark Liao},
+      year={2024},
+      eprint={2402.13616},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV}
+}
+```

docs/CONTRIBUTING.md

@@ -0,0 +1,44 @@
+# Contributing to YOLO
+
+Thank you for your interest in contributing to this project! We value your contributions and want to make the process as easy and enjoyable as possible. Below you will find the guidelines for contributing.
+
+## Quick Links
+- [Main README](../README.md)
+- [License](../LICENSE)
+- [Issue Tracker](https://github.com/WongKinYiu/yolov9mit/issues)
+- [Pull Requests](https://github.com/WongKinYiu/yolov9mit/pulls)
+
+## Testing and Formatting
+We strive to maintain a high standard of quality in our codebase:
+- **Testing:** We use `pytest` for testing. Please add tests for new code you create.
+- **Formatting:** Our code follows a consistent style enforced by `isort` for imports sorting and `black` for code formatting. Run these tools to format your code before submitting a pull request.
+
+## GitHub Actions
+We utilize GitHub Actions for continuous integration. When you submit a pull request, automated tests and formatting checks will run. Ensure that these checks pass for your pull request to be accepted.
+
+## How to Contribute
+
+### Proposing Enhancements
+For feature requests or improvements, open an issue with:
+- A clear title and description.
+- Explain why this enhancement would be useful.
+- Considerations or potential implementation details.
+
+## Pull Request Checklist
+Before sending your pull request, always check the following:
+- The code follows the [Python style guide](https://www.python.org/dev/peps/pep-0008/).
+- Code and files are well organized.
+- All tests pass.
+- New code is covered by tests.
+- We would be very happy if [gitmoji😆](https://www.npmjs.com/package/gitmoji-cli) could be used to assist the commit message💬!
+
+## Code Review Process
+Once you submit a PR, maintainers will review your work, suggest changes if necessary, and merge it once it’s approved.
+
+---
+
+Your contributions are greatly appreciated and vital to the project's success!
+
+Please feel free to contact [[email protected]](mailto:[email protected])!
+
+

docs/HOWTO.md

@@ -0,0 +1,90 @@
+# How To modified YOLO
+
+To facilitate easy customization of the YOLO model, we've structured the codebase to allow for changes through configuration files and minimal code adjustments. This guide will walk you through the steps to customize various components of the model including the architecture, blocks, data loaders, and loss functions.
+
+## Custom Model Architecture
+
+You can change the model architecture simply by modifying the YAML configuration file. Here's how:
+
+1. **Modify Architecture in Config:**
+   
+   Navigate to your model's configuration file (typically formate like `yolo/config/model/v9-c.yaml`).
+   - Adjust the architecture settings under the `architecture` section. Ensure that every module you reference exists in `module.py`, or refer to the next section on how to add new modules.
+
+    ```yaml
+    model:
+      foo:
+        - ADown:
+            args: {out_channels: 256}
+        - RepNCSPELAN:
+            source: -2
+            args: {out_channels: 512, part_channels: 256}
+            tags: B4
+      bar:
+        - Concat:
+            source: [-2, B4]
+    ```
+
+   `tags`: Use this to labels any module you want, and could be the module source.
+
+   `source`: Set this to the index of the module output you wish to use as input; default is `-1` which refers to the last module's output. Capable tags, relative position, absolute position
+
+   `args`: A dictionary used to initialize parameters for convolutional or bottleneck layers.
+
+   `output`: Whether to serve as the output of the model.
+
+## Custom Block
+
+To add or modify a block in the model:
+
+1. **Create a New Module:**
+
+   Define a new class in `module.py` that inherits from `nn.Module`.
+   
+   The constructor should accept `in_channels` as a parameter. Make sure to calculate `out_channels` based on your model's requirements or configure it through the YAML file using `args`.
+
+    ```python
+    class CustomBlock(nn.Module):
+        def __init__(self, in_channels, out_channels, **kwargs):
+            super().__init__()
+            self.module = # conv, bool, ...
+        def forward(self, x):
+            return self.module(x) 
+    ```
+
+2. **Reference in Config:**
+   ```yaml
+    ...
+    - CustomBlock:
+        args: {out_channels: int, etc: ...}
+        ...
+    ...
+   ```
+
+
+## Custom Data Augmentation
+
+Custom transformations should be designed to accept an image and its bounding boxes, and return them after applying the desired changes. Here’s how you can define such a transformation:
+
+
+1. **Define Dataset:**
+
+    Your class must have a `__call__` method that takes a PIL image and its corresponding bounding boxes as input, and returns them after processing.
+
+
+   ```python
+    class CustomTransform:
+        def __init__(self, prob=0.5):
+            self.prob = prob
+
+        def __call__(self, image, boxes):
+            return image, boxes
+   ```
+2. **Update CustomTransform in Config:**
+
+    Specify your custom transformation in a YAML config `yolo/config/data/augment.yaml`. For examples:
+    ```yaml
+    Mosaic: 1
+    # ... (Other Transform)
+    CustomTransform: 0.5
+    ```

docs/MODELS.md

@@ -0,0 +1,31 @@
+# YOLO Model Zoo
+
+Welcome to the YOLOv9 Model Zoo! Here, you will find a variety of pre-trained models tailored to different use cases and performance needs. Each model comes with detailed information about its training regime, performance metrics, and usage instructions.
+
+## Standard Models
+
+These models are trained on common datasets like COCO and provide a balance between speed and accuracy.
+
+
+| Model | Test Size | AP<sup>val</sup> | AP<sub>50</sub><sup>val</sup> | AP<sub>75</sub><sup>val</sup> | Param. | FLOPs |
+| :-- | :-: | :-: | :-: | :-: | :-: | :-: |
+| [**YOLOv9-T**]() | 640 | **38.3%** | **53.1%** | **41.3%** | **2.0M** | **7.7G** |
+| [**YOLOv9-S**]() | 640 | **46.8%** | **63.4%** | **50.7%** | **7.1M** | **26.4G** |
+| [**YOLOv9-M**]() | 640 | **51.4%** | **68.1%** | **56.1%** | **20.0M** | **76.3G** |
+| [**YOLOv9-C**]() | 640 | **53.0%** | **70.2%** | **57.8%** | **25.3M** | **102.1G** |
+| [**YOLOv9-E**]() | 640 | **55.6%** | **72.8%** | **60.6%** | **57.3M** | **189.0G** |
+|  |  |  |  |  |  |  |
+| [**YOLOv7**]() | 640 | **51.4%** | **69.7%** | **55.9%** |
+| [**YOLOv7-X**]() | 640 | **53.1%** | **71.2%** | **57.8%** |
+| [**YOLOv7-W6**]() | 1280 | **54.9%** | **72.6%** | **60.1%** |
+| [**YOLOv7-E6**]() | 1280 | **56.0%** | **73.5%** | **61.2%** |
+| [**YOLOv7-D6**]() | 1280 | **56.6%** | **74.0%** | **61.8%** |
+| [**YOLOv7-E6E**]() | 1280 | **56.8%** | **74.4%** | **62.1%** |
+
+## Download and Usage Instructions
+
+To use these models, download them from the links provided and use the following command to run detection:
+
+```bash
+$yolo detect weights=path/to/model.pt img=640 conf=0.25 source=your_image.jpg
+```

examples/example_train.py

@@ -9,29 +9,25 @@ 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.drawer import draw_model
 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)
-    draw_model(model=model)
     # 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/model/module.py

@@ -25,7 +25,7 @@ class Conv(nn.Module):
         super().__init__()
         kwargs.setdefault("padding", auto_pad(kernel_size, **kwargs))
         self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, bias=False, **kwargs)
-        self.bn = nn.BatchNorm2d(out_channels)
+        self.bn = nn.BatchNorm2d(out_channels, eps=1e-3, momentum=3e-2)
         self.act = get_activation(activation)
 
     def forward(self, x: Tensor) -> Tensor:
@@ -69,6 +69,9 @@ class Detection(nn.Module):
             Conv(in_channels, class_neck, 3), Conv(class_neck, class_neck, 3), nn.Conv2d(class_neck, num_classes, 1)
         )
 
+        self.anchor_conv[-1].bias.data.fill_(1.0)
+        self.class_conv[-1].bias.data.fill_(-10)
+
     def forward(self, x: List[Tensor]) -> List[Tensor]:
         anchor_x = self.anchor_conv(x)
         class_x = self.class_conv(x)

yolo/model/yolo.py

@@ -7,6 +7,7 @@ from omegaconf import ListConfig, OmegaConf
 from yolo.config.config import Config, Model, YOLOLayer
 from yolo.tools.layer_helper import get_layer_map
 from yolo.tools.log_helper import log_model
+from yolo.utils.drawer import draw_model
 
 
 class YOLO(nn.Module):
@@ -24,8 +25,6 @@ class YOLO(nn.Module):
         self.layer_map = get_layer_map()  # Get the map Dict[str: Module]
         self.model: List[YOLOLayer] = nn.ModuleList()
         self.build_model(model_cfg.model)
-        # TODO: Move to other position
-        log_model(self.model)
 
     def build_model(self, model_arch: Dict[str, List[Dict[str, Dict[str, Dict]]]]):
         self.layer_index = {}
@@ -126,4 +125,6 @@ def get_model(cfg: Config) -> YOLO:
     OmegaConf.set_struct(cfg.model, False)
     model = YOLO(cfg.model, cfg.hyper.data.class_num)
     logger.info("✅ Success load model")
+    log_model(model.model)
+    # draw_model(model=model)
     return model

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", total=num_epochs)
+        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]):
@@ -63,21 +66,25 @@ class CustomProgress:
             for loss_name, loss_value in loss_dict.items():
                 self.wandb.log({f"Loss/{loss_name}": loss_value})
 
-        loss_str = "Loss"
+        loss_str = "| -.-- |"
         for loss_name, loss_val in loss_dict.items():
-            loss_str += f" {loss_name[:-4]}: {loss_val:.2f} |"
+            loss_str += f" {loss_val:2.2f} |"
 
         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):
 
-                epoch_loss = self.train_one_epoch(dataloader, self.progress)
-                self.progress.one_epoch()
+                self.progress.start_one_epoch(len(dataloader), self.optimizer, epoch)
+                epoch_loss = self.train_one_epoch(dataloader)
+                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

@@ -7,6 +7,7 @@ import numpy as np
 import torch
 from loguru import logger
 from PIL import Image
+from rich.progress import track
 from torch.utils.data import DataLoader, Dataset
 from torchvision.transforms import functional as TF
 from tqdm.rich import tqdm
@@ -74,7 +75,7 @@ class YoloDataset(Dataset):
 
         data = []
         valid_inputs = 0
-        for image_name in tqdm(images_list, desc="Filtering data"):
+        for image_name in track(images_list, description="Filtering data"):
             if not image_name.lower().endswith((".jpg", ".jpeg", ".png")):
                 continue
             image_id, _ = path.splitext(image_name)
@@ -159,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,
         )