✨ [New] use lightning framework to training!

yolo/__init__.py

@@ -2,18 +2,22 @@ from yolo.config.config import Config, NMSConfig
 from yolo.model.yolo import create_model
 from yolo.tools.data_loader import AugmentationComposer, create_dataloader
 from yolo.tools.drawer import draw_bboxes
-from yolo.tools.solver import ModelTester, ModelTrainer, ModelValidator
+from yolo.tools.solver import TrainModel
 from yolo.utils.bounding_box_utils import Anc2Box, Vec2Box, bbox_nms, create_converter
 from yolo.utils.deploy_utils import FastModelLoader
-from yolo.utils.logging_utils import ProgressLogger, custom_logger
+from yolo.utils.logging_utils import (
+    ImageLogger,
+    YOLORichModelSummary,
+    YOLORichProgressBar,
+)
 from yolo.utils.model_utils import PostProccess
 
 all = [
     "create_model",
     "Config",
-    "ProgressLogger",
+    "YOLORichProgressBar",
     "NMSConfig",
-    "custom_logger",
+    "YOLORichModelSummary",
     "validate_log_directory",
     "draw_bboxes",
     "Vec2Box",
@@ -21,10 +25,9 @@ all = [
     "bbox_nms",
     "create_converter",
     "AugmentationComposer",
+    "ImageLogger",
     "create_dataloader",
     "FastModelLoader",
-    "ModelTester",
-    "ModelTrainer",
-    "ModelValidator",
+    "TrainModel",
     "PostProccess",
 ]

yolo/lazy.py

@@ -2,41 +2,36 @@ import sys
 from pathlib import Path
 
 import hydra
+from lightning import Trainer
 
 project_root = Path(__file__).resolve().parent.parent
 sys.path.append(str(project_root))
 
 from yolo.config.config import Config
-from yolo.model.yolo import create_model
-from yolo.tools.data_loader import create_dataloader
-from yolo.tools.solver import ModelTester, ModelTrainer, ModelValidator
-from yolo.utils.bounding_box_utils import create_converter
-from yolo.utils.deploy_utils import FastModelLoader
-from yolo.utils.logging_utils import ProgressLogger
-from yolo.utils.model_utils import get_device
+from yolo.tools.solver import TrainModel, ValidateModel
+from yolo.utils.logging_utils import setup
 
 
 @hydra.main(config_path="config", config_name="config", version_base=None)
 def main(cfg: Config):
-    progress = ProgressLogger(cfg, exp_name=cfg.name)
-    device, use_ddp = get_device(cfg.device)
-    dataloader = create_dataloader(cfg.task.data, cfg.dataset, cfg.task.task, use_ddp)
-    if getattr(cfg.task, "fast_inference", False):
-        model = FastModelLoader(cfg).load_model(device)
-    else:
-        model = create_model(cfg.model, class_num=cfg.dataset.class_num, weight_path=cfg.weight)
-        model = model.to(device)
-
-    converter = create_converter(cfg.model.name, model, cfg.model.anchor, cfg.image_size, device)
-
-    if cfg.task.task == "train":
-        solver = ModelTrainer(cfg, model, converter, progress, device, use_ddp)
-    if cfg.task.task == "validation":
-        solver = ModelValidator(cfg.task, cfg.dataset, model, converter, progress, device)
-    if cfg.task.task == "inference":
-        solver = ModelTester(cfg, model, converter, progress, device)
-    progress.start()
-    solver.solve(dataloader)
+    callbacks, loggers = setup(cfg)
+
+    trainer = Trainer(
+        accelerator="cuda",
+        max_epochs=getattr(cfg.task, "epoch", None),
+        precision="16-mixed",
+        callbacks=callbacks,
+        logger=loggers,
+        log_every_n_steps=1,
+    )
+
+    match cfg.task.task:
+        case "train":
+            model = TrainModel(cfg)
+            trainer.fit(model)
+        case "validation":
+            model = ValidateModel(cfg)
+            trainer.validate(model)
 
 
 if __name__ == "__main__":

yolo/tools/solver.py

@@ -1,267 +1,89 @@
-import contextlib
-import io
-import json
-import os
-import time
-from collections import defaultdict
-from pathlib import Path
-from typing import Dict, Optional
+from lightning import LightningModule
+from torchmetrics.detection import MeanAveragePrecision
 
-import torch
-from pycocotools.coco import COCO
-from torch import Tensor, distributed
-from torch.cuda.amp import GradScaler, autocast
-from torch.nn.parallel import DistributedDataParallel as DDP
-from torch.utils.data import DataLoader
-
-from yolo.config.config import Config, DatasetConfig, TrainConfig, ValidationConfig
-from yolo.model.yolo import YOLO
-from yolo.tools.data_loader import StreamDataLoader, create_dataloader
-from yolo.tools.drawer import draw_bboxes, draw_model
+from yolo.config.config import Config
+from yolo.model.yolo import create_model
+from yolo.tools.data_loader import create_dataloader
 from yolo.tools.loss_functions import create_loss_function
-from yolo.utils.bounding_box_utils import Vec2Box, calculate_map
-from yolo.utils.dataset_utils import locate_label_paths
-from yolo.utils.logger import logger
-from yolo.utils.logging_utils import ProgressLogger, log_model_structure
-from yolo.utils.model_utils import (
-    ExponentialMovingAverage,
-    PostProccess,
-    collect_prediction,
-    create_optimizer,
-    create_scheduler,
-    predicts_to_json,
-)
-from yolo.utils.solver_utils import calculate_ap
-
+from yolo.utils.bounding_box_utils import create_converter, to_metrics_format
+from yolo.utils.model_utils import PostProccess, create_optimizer, create_scheduler
 
-class ModelTrainer:
-    def __init__(self, cfg: Config, model: YOLO, vec2box: Vec2Box, progress: ProgressLogger, device, use_ddp: bool):
-        train_cfg: TrainConfig = cfg.task
-        self.model = model if not use_ddp else DDP(model, device_ids=[device])
-        self.use_ddp = use_ddp
-        self.vec2box = vec2box
-        self.device = device
-        self.optimizer = create_optimizer(model, train_cfg.optimizer)
-        self.scheduler = create_scheduler(self.optimizer, train_cfg.scheduler)
-        self.loss_fn = create_loss_function(cfg, vec2box)
-        self.progress = progress
-        self.num_epochs = cfg.task.epoch
-        self.mAPs_dict = defaultdict(list)
 
-        self.weights_dir = self.progress.save_path / "weights"
-        self.weights_dir.mkdir(exist_ok=True)
+class BaseModel(LightningModule):
+    def __init__(self, cfg: Config):
+        super().__init__()
+        self.model = create_model(cfg.model, class_num=cfg.dataset.class_num, weight_path=cfg.weight)
 
-        if not progress.quite_mode:
-            log_model_structure(model.model)
-            draw_model(model=model)
+    def forward(self, x):
+        return self.model(x)
 
-        self.validation_dataloader = create_dataloader(
-            cfg.task.validation.data, cfg.dataset, cfg.task.validation.task, use_ddp
-        )
-        self.validator = ModelValidator(cfg.task.validation, cfg.dataset, model, vec2box, progress, device)
 
-        if getattr(train_cfg.ema, "enabled", False):
-            self.ema = ExponentialMovingAverage(model, decay=train_cfg.ema.decay)
+class ValidateModel(BaseModel):
+    def __init__(self, cfg: Config):
+        super().__init__(cfg)
+        self.cfg = cfg
+        if self.cfg.task.task == "validation":
+            self.validation_cfg = self.cfg.task
         else:
-            self.ema = None
-        self.scaler = GradScaler()
-
-    def train_one_batch(self, images: Tensor, targets: Tensor):
-        images, targets = images.to(self.device), targets.to(self.device)
-        self.optimizer.zero_grad()
-
-        with autocast():
-            predicts = self.model(images)
-            aux_predicts = self.vec2box(predicts["AUX"])
-            main_predicts = self.vec2box(predicts["Main"])
-            loss, loss_item = self.loss_fn(aux_predicts, main_predicts, targets)
-
-        self.scaler.scale(loss).backward()
-        self.scaler.unscale_(self.optimizer)
-        torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=10.0)
-        self.scaler.step(self.optimizer)
-        self.scaler.update()
-
-        return loss_item
-
-    def train_one_epoch(self, dataloader):
-        self.model.train()
-        total_loss = defaultdict(float)
-        total_samples = 0
-        self.optimizer.next_epoch(len(dataloader))
-        for batch_size, images, targets, *_ in dataloader:
-            self.optimizer.next_batch()
-            loss_each = self.train_one_batch(images, targets)
-
-            for loss_name, loss_val in loss_each.items():
-                if self.use_ddp:  # collecting loss for each batch
-                    distributed.all_reduce(loss_val, op=distributed.ReduceOp.AVG)
-                total_loss[loss_name] += loss_val.item() * batch_size
-            total_samples += batch_size
-            self.progress.one_batch(loss_each)
-
-        for loss_val in total_loss.values():
-            loss_val /= total_samples
-
-        if self.scheduler:
-            self.scheduler.step()
-
-        return total_loss
-
-    def save_checkpoint(self, epoch_idx: int, file_name: Optional[str] = None):
-        file_name = file_name or f"E{epoch_idx:03d}.pt"
-        file_path = self.weights_dir / file_name
-
-        checkpoint = {
-            "epoch": epoch_idx,
-            "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()
-
-        logger.info(f"💾 success save at {file_path}")
-        torch.save(checkpoint, file_path)
-
-    def good_epoch(self, mAPs: Dict[str, Tensor]) -> bool:
-        save_flag = True
-        for mAP_key, mAP_val in mAPs.items():
-            self.mAPs_dict[mAP_key].append(mAP_val)
-            if mAP_val < max(self.mAPs_dict[mAP_key]):
-                save_flag = False
-        return save_flag
-
-    def solve(self, dataloader: DataLoader):
-        logger.info("🚄 Start Training!")
-        num_epochs = self.num_epochs
-
-        self.progress.start_train(num_epochs)
-        for epoch_idx in range(num_epochs):
-            if self.use_ddp:
-                dataloader.sampler.set_epoch(epoch_idx)
+            self.validation_cfg = self.cfg.task.validation
+        self.metric = MeanAveragePrecision(iou_type="bbox", box_format="xyxy")
 
-            self.progress.start_one_epoch(len(dataloader), "Train", self.optimizer, epoch_idx)
-            epoch_loss = self.train_one_epoch(dataloader)
-            self.progress.finish_one_epoch(epoch_loss, epoch_idx=epoch_idx)
-
-            mAPs = self.validator.solve(self.validation_dataloader, epoch_idx=epoch_idx)
-            if mAPs is not None and self.good_epoch(mAPs):
-                self.save_checkpoint(epoch_idx=epoch_idx)
-            # TODO: save model if result are better than before
-        self.progress.finish_train()
-
-
-class ModelTester:
-    def __init__(self, cfg: Config, model: YOLO, vec2box: Vec2Box, progress: ProgressLogger, device):
-        self.model = model
-        self.device = device
-        self.progress = progress
-
-        self.post_proccess = PostProccess(vec2box, cfg.task.nms)
-        self.save_path = progress.save_path / "images"
-        os.makedirs(self.save_path, exist_ok=True)
-        self.save_predict = getattr(cfg.task, "save_predict", None)
-        self.idx2label = cfg.dataset.class_list
-
-    def solve(self, dataloader: StreamDataLoader):
-        logger.info("👀 Start Inference!")
-        if isinstance(self.model, torch.nn.Module):
-            self.model.eval()
-
-        if dataloader.is_stream:
-            import cv2
-            import numpy as np
-
-            last_time = time.time()
-        try:
-            for idx, (images, rev_tensor, origin_frame) in enumerate(dataloader):
-                images = images.to(self.device)
-                rev_tensor = rev_tensor.to(self.device)
-                with torch.no_grad():
-                    predicts = self.model(images)
-                    predicts = self.post_proccess(predicts, rev_tensor)
-                img = draw_bboxes(origin_frame, predicts, idx2label=self.idx2label)
-
-                if dataloader.is_stream:
-                    img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
-                    fps = 1 / (time.time() - last_time)
-                    cv2.putText(img, f"FPS: {fps:.2f}", (0, 15), 0, 0.5, (100, 255, 0), 1, cv2.LINE_AA)
-                    last_time = time.time()
-                    cv2.imshow("Prediction", img)
-                    if cv2.waitKey(1) & 0xFF == ord("q"):
-                        break
-                    if not self.save_predict:
-                        continue
-                if self.save_predict != False:
-                    save_image_path = self.save_path / f"frame{idx:03d}.png"
-                    img.save(save_image_path)
-                    logger.info(f"💾 Saved visualize image at {save_image_path}")
-
-        except (KeyboardInterrupt, Exception) as e:
-            dataloader.stop_event.set()
-            dataloader.stop()
-            if isinstance(e, KeyboardInterrupt):
-                logger.error("User Keyboard Interrupt")
-            else:
-                raise e
-        dataloader.stop()
-
-
-class ModelValidator:
-    def __init__(
-        self,
-        validation_cfg: ValidationConfig,
-        dataset_cfg: DatasetConfig,
-        model: YOLO,
-        vec2box: Vec2Box,
-        progress: ProgressLogger,
-        device,
-    ):
-        self.model = model
-        self.device = device
-        self.progress = progress
-
-        self.post_proccess = PostProccess(vec2box, validation_cfg.nms)
-        self.json_path = self.progress.save_path / "predict.json"
-
-        with contextlib.redirect_stdout(io.StringIO()):
-            # TODO: load with config file
-            json_path, _ = locate_label_paths(Path(dataset_cfg.path), dataset_cfg.get("validation", "val"))
-            if json_path:
-                self.coco_gt = COCO(json_path)
-
-    def solve(self, dataloader, epoch_idx=1):
-        # logger.info("🧪 Start Validation!")
-        self.model.eval()
-        predict_json, mAPs = [], defaultdict(list)
-        self.progress.start_one_epoch(len(dataloader), task="Validate")
-        for batch_size, images, targets, rev_tensor, img_paths in dataloader:
-            images, targets, rev_tensor = images.to(self.device), targets.to(self.device), rev_tensor.to(self.device)
-            with torch.no_grad():
-                predicts = self.model(images)
-                predicts = self.post_proccess(predicts)
-                for idx, predict in enumerate(predicts):
-                    mAP = calculate_map(predict, targets[idx])
-                    for mAP_key, mAP_val in mAP.items():
-                        mAPs[mAP_key].append(mAP_val)
-
-            avg_mAPs = {key: 100 * torch.mean(torch.stack(val)) for key, val in mAPs.items()}
-            self.progress.one_batch(avg_mAPs)
+    def setup(self, stage):
+        self.vec2box = create_converter(
+            self.cfg.model.name, self.model, self.cfg.model.anchor, self.cfg.image_size, self.device
+        )
+        self.post_proccess = PostProccess(self.vec2box, self.validation_cfg.nms)
 
-            predict_json.extend(predicts_to_json(img_paths, predicts, rev_tensor))
-        self.progress.finish_one_epoch(avg_mAPs, epoch_idx=epoch_idx)
-        self.progress.visualize_image(images, targets, predicts, epoch_idx=epoch_idx)
+    def val_dataloader(self):
+        return create_dataloader(self.validation_cfg.data, self.cfg.dataset, self.validation_cfg.task)
 
-        with open(self.json_path, "w") as f:
-            predict_json = collect_prediction(predict_json, self.progress.local_rank)
-            if self.progress.local_rank != 0:
-                return
-            json.dump(predict_json, f)
-        if hasattr(self, "coco_gt"):
-            self.progress.start_pycocotools()
-            result = calculate_ap(self.coco_gt, predict_json)
-            self.progress.finish_pycocotools(result, epoch_idx)
+    def validation_step(self, batch, batch_idx):
+        batch_size, images, targets, rev_tensor, img_paths = batch
+        predicts = self.post_proccess(self(images))
+        batch_metrics = self.metric(
+            [to_metrics_format(predict) for predict in predicts], [to_metrics_format(target) for target in targets]
+        )
 
-        return avg_mAPs
+        self.log_dict(
+            {
+                "map": batch_metrics["map"],
+                "map_50": batch_metrics["map_50"],
+            },
+            on_step=True,
+            prog_bar=True,
+            logger=False,
+            batch_size=batch_size,
+        )
+        return predicts
+
+    def on_validation_epoch_end(self):
+        epoch_metrics = self.metric.compute()
+        del epoch_metrics["classes"]
+        self.log_dict(epoch_metrics, on_epoch=True, prog_bar=True, logger=True)
+
+
+class TrainModel(ValidateModel):
+    def __init__(self, cfg: Config):
+        super().__init__(cfg)
+        self.cfg = cfg
+
+    def setup(self, stage):
+        super().setup(stage)
+        self.loss_fn = create_loss_function(self.cfg, self.vec2box)
+
+    def train_dataloader(self):
+        return create_dataloader(self.cfg.task.data, self.cfg.dataset, self.cfg.task.task)
+
+    def training_step(self, batch, batch_idx):
+        batch_size, images, targets, *_ = batch
+        predicts = self(images)
+        aux_predicts = self.vec2box(predicts["AUX"])
+        main_predicts = self.vec2box(predicts["Main"])
+        loss, loss_item = self.loss_fn(aux_predicts, main_predicts, targets)
+        self.log_dict(loss_item, on_step=True, on_epoch=True, prog_bar=True, logger=True, batch_size=batch_size)
+        return loss * batch_size
+
+    def configure_optimizers(self):
+        optimizer = create_optimizer(self.model, self.cfg.task.optimizer)
+        scheduler = create_scheduler(optimizer, self.cfg.task.scheduler)
+        return [optimizer], [scheduler]

yolo/utils/bounding_box_utils.py

@@ -446,3 +446,10 @@ def calculate_map(predictions, ground_truths, iou_thresholds=arange(0.5, 1, 0.05
         "mAP.5:.95": torch.mean(torch.stack(aps)),
     }
     return mAP
+
+
+def to_metrics_format(prediction: Tensor) -> Dict[str, Union[float, Tensor]]:
+    bbox = {"boxes": prediction[:, 1:5], "labels": prediction[:, 0].int()}
+    if prediction.size(1) == 6:
+        bbox["scores"] = prediction[:, 5]
+    return bbox

yolo/utils/logging_utils.py

@@ -11,9 +11,7 @@ Example:
     custom_logger()
 """
 
-import os
-import random
-import sys
+import logging
 from collections import deque
 from logging import FileHandler
 from pathlib import Path
@@ -22,39 +20,29 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 import numpy as np
 import torch
 import wandb
-import wandb.errors.term
+from lightning import LightningModule, Trainer, seed_everything
+from lightning.pytorch.callbacks import Callback, RichModelSummary, RichProgressBar
+from lightning.pytorch.callbacks.progress.rich_progress import CustomProgress
+from lightning.pytorch.loggers import WandbLogger
 from omegaconf import ListConfig
+from rich import reconfigure
 from rich.console import Console, Group
-from rich.progress import (
-    BarColumn,
-    Progress,
-    SpinnerColumn,
-    TextColumn,
-    TimeRemainingColumn,
-)
+from rich.logging import RichHandler
 from rich.table import Table
+from rich.text import Text
 from torch import Tensor
 from torch.nn import ModuleList
-from torch.optim import Optimizer
-from torchvision.transforms.functional import pil_to_tensor
+from typing_extensions import override
 
 from yolo.config.config import Config, YOLOLayer
 from yolo.model.yolo import YOLO
-from yolo.tools.drawer import draw_bboxes
 from yolo.utils.logger import logger
 from yolo.utils.solver_utils import make_ap_table
 
 
-def custom_logger(quite: bool = False):
-    if quite:
-        logger.removeHandler("YOLO_logger")
-
-
 # TODO: should be moved to correct position
 def set_seed(seed):
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
+    seed_everything(seed)
     if torch.cuda.is_available():
         torch.cuda.manual_seed(seed)
         torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
@@ -62,189 +50,211 @@ def set_seed(seed):
     torch.backends.cudnn.benchmark = False
 
 
-class ProgressLogger(Progress):
-    def __init__(self, cfg: Config, exp_name: str, *args, **kwargs):
-        set_seed(cfg.lucky_number)
-        self.local_rank = int(os.getenv("LOCAL_RANK", "0"))
-        self.quite_mode = self.local_rank or getattr(cfg, "quite", False)
-        custom_logger(self.quite_mode)
-        self.save_path = validate_log_directory(cfg, exp_name=cfg.name)
-
-        progress_bar = (
-            SpinnerColumn(),
-            TextColumn("[progress.description]{task.description}"),
-            BarColumn(bar_width=None),
-            TextColumn("{task.completed:.0f}/{task.total:.0f}"),
-            TimeRemainingColumn(),
-        )
-        self.ap_table = Table()
-        # TODO: load maxlen by config files
-        self.ap_past_list = deque(maxlen=5)
-        self.last_result = 0
-        super().__init__(*args, *progress_bar, **kwargs)
-
-        self.use_wandb = cfg.use_wandb
-        if self.use_wandb and self.local_rank == 0:
-            wandb.errors.term._log = custom_wandb_log
-            self.wandb = wandb.init(
-                project="YOLO", resume="allow", mode="online", dir=self.save_path, id=None, name=exp_name
-            )
+class YOLOCustomProgress(CustomProgress):
+    def get_renderable(self):
+        renderable = Group(*self.get_renderables())
+        if hasattr(self, "table"):
+            renderable = Group(*self.get_renderables(), self.table)
+        return renderable
 
-        self.use_tensorboard = cfg.use_tensorboard
-        if self.use_tensorboard and self.local_rank == 0:
-            from torch.utils.tensorboard import SummaryWriter
 
-            self.tb_writer = SummaryWriter(log_dir=self.save_path / "tensorboard")
-            logger.info(f"📍 Enable TensorBoard locally at <blue><u>http://localhost:6006</></>")
+class YOLORichProgressBar(RichProgressBar):
+    @override
+    def _init_progress(self, trainer: "Trainer") -> None:
+        if self.is_enabled and (self.progress is None or self._progress_stopped):
+            self._reset_progress_bar_ids()
+            reconfigure(**self._console_kwargs)
+            self._console = Console()
+            self._console.clear_live()
+            self.progress = YOLOCustomProgress(
+                *self.configure_columns(trainer),
+                auto_refresh=False,
+                disable=self.is_disabled,
+                console=self._console,
+            )
+            self.progress.start()
 
-    def rank_check(logging_function):
-        def wrapper(self, *args, **kwargs):
-            if getattr(self, "local_rank", 0) != 0:
-                return
-            return logging_function(self, *args, **kwargs)
+            self._progress_stopped = False
 
-        return wrapper
+            self.max_result = 0
+            self.past_results = deque(maxlen=5)
+            self.progress.table = Table()
 
-    def get_renderable(self):
-        renderable = Group(*self.get_renderables(), self.ap_table)
-        return renderable
+    @override
+    def _get_train_description(self, current_epoch: int) -> str:
+        return Text("[cyan]Train [white]|")
 
-    @rank_check
-    def start_train(self, num_epochs: int):
-        self.task_epoch = self.add_task(f"[cyan]Start Training {num_epochs} epochs", total=num_epochs)
-        self.update(self.task_epoch, advance=-0.5)
-
-    @rank_check
-    def start_one_epoch(
-        self, num_batches: int, task: str = "Train", optimizer: Optimizer = None, epoch_idx: int = None
-    ):
-        self.num_batches = num_batches
-        self.task = task
-        if hasattr(self, "task_epoch"):
-            self.update(self.task_epoch, description=f"[cyan] Preparing Data")
-
-        if optimizer is not None:
-            lr_values = [params["lr"] for params in optimizer.param_groups]
-            lr_names = ["Learning Rate/bias", "Learning Rate/norm", "Learning Rate/conv"]
-            if self.use_wandb:
-                for lr_name, lr_value in zip(lr_names, lr_values):
-                    self.wandb.log({lr_name: lr_value}, step=epoch_idx)
-
-            if self.use_tensorboard:
-                for lr_name, lr_value in zip(lr_names, lr_values):
-                    self.tb_writer.add_scalar(lr_name, lr_value, global_step=epoch_idx)
-
-        self.batch_task = self.add_task(f"[green] Phase: {task}", total=num_batches)
-
-    @rank_check
-    def one_batch(self, batch_info: Dict[str, Tensor] = None):
-        epoch_descript = "[cyan]" + self.task + "[white] |"
-        batch_descript = "|"
-        if self.task == "Train":
-            self.update(self.task_epoch, advance=1 / self.num_batches)
-        for info_name, info_val in batch_info.items():
-            epoch_descript += f"{info_name: ^9}|"
-            batch_descript += f"   {info_val:2.2f}  |"
-        self.update(self.batch_task, advance=1, description=f"[green]{self.task} [white]{batch_descript}")
-        if hasattr(self, "task_epoch"):
-            self.update(self.task_epoch, description=epoch_descript)
-
-    @rank_check
-    def finish_one_epoch(self, batch_info: Dict[str, Any] = None, epoch_idx: int = -1):
-        if self.task == "Train":
-            prefix = "Loss"
-        elif self.task == "Validate":
-            prefix = "Metrics"
-        batch_info = {f"{prefix}/{key}": value for key, value in batch_info.items()}
-        if self.use_wandb:
-            self.wandb.log(batch_info, step=epoch_idx)
-        if self.use_tensorboard:
-            for key, value in batch_info.items():
-                self.tb_writer.add_scalar(key, value, epoch_idx)
-
-        self.remove_task(self.batch_task)
-
-    @rank_check
-    def visualize_image(
-        self,
-        images: Optional[Tensor] = None,
-        ground_truth: Optional[Tensor] = None,
-        prediction: Optional[Union[List[Tensor], Tensor]] = None,
-        epoch_idx: int = 0,
-    ) -> None:
-        """
-        Upload the ground truth bounding boxes, predicted bounding boxes, and the original image to wandb or TensorBoard.
-
-        Args:
-            images (Optional[Tensor]): Tensor of images with shape (BZ, 3, 640, 640).
-            ground_truth (Optional[Tensor]): Ground truth bounding boxes with shape (BZ, N, 5) or (N, 5). Defaults to None.
-            prediction (prediction: Optional[Union[List[Tensor], Tensor]]): List of predicted bounding boxes with shape (N, 6) or (N, 6). Defaults to None.
-            epoch_idx (int): Current epoch index. Defaults to 0.
-        """
-        if images is not None:
-            images = images[0] if images.ndim == 4 else images
-            if self.use_wandb:
-                wandb.log({"Input Image": wandb.Image(images)}, step=epoch_idx)
-            if self.use_tensorboard:
-                self.tb_writer.add_image("Media/Input Image", images, 1)
-
-        if ground_truth is not None:
-            gt_boxes = ground_truth[0] if ground_truth.ndim == 3 else ground_truth
-            if self.use_wandb:
-                wandb.log(
-                    {"Ground Truth": wandb.Image(images, boxes={"predictions": {"box_data": log_bbox(gt_boxes)}})},
-                    step=epoch_idx,
-                )
-            if self.use_tensorboard:
-                self.tb_writer.add_image("Media/Ground Truth", pil_to_tensor(draw_bboxes(images, gt_boxes)), epoch_idx)
-
-        if prediction is not None:
-            pred_boxes = prediction[0] if isinstance(prediction, list) else prediction
-            if self.use_wandb:
-                wandb.log(
-                    {"Prediction": wandb.Image(images, boxes={"predictions": {"box_data": log_bbox(pred_boxes)}})},
-                    step=epoch_idx,
-                )
-            if self.use_tensorboard:
-                self.tb_writer.add_image("Media/Prediction", pil_to_tensor(draw_bboxes(images, pred_boxes)), epoch_idx)
-
-    @rank_check
-    def start_pycocotools(self):
-        self.batch_task = self.add_task("[green]Run pycocotools", total=1)
-
-    @rank_check
-    def finish_pycocotools(self, result, epoch_idx=-1):
-        ap_table, ap_main = make_ap_table(result * 100, self.ap_past_list, self.last_result, epoch_idx)
-        self.last_result = np.maximum(result, self.last_result)
-        self.ap_past_list.append((epoch_idx, ap_main))
-        self.ap_table = ap_table
-
-        if self.use_wandb:
-            self.wandb.log({"PyCOCO/AP @ .5:.95": ap_main[2], "PyCOCO/AP @ .5": ap_main[5]})
-        if self.use_tensorboard:
-            # TODO: waiting torch bugs fix, https://github.com/pytorch/pytorch/issues/32651
-            self.tb_writer.add_scalar("PyCOCO/AP @ .5:.95", ap_main[2], epoch_idx)
-            self.tb_writer.add_scalar("PyCOCO/AP @ .5", ap_main[5], epoch_idx)
-
-        self.update(self.batch_task, advance=1)
+    @override
+    def on_train_start(self, trainer, pl_module):
+        self._init_progress(trainer)
+        num_epochs = trainer.max_epochs - 1
+        self.task_epoch = self._add_task(
+            total_batches=num_epochs,
+            description=f"[cyan]Start Training {num_epochs} epochs",
+        )
+        self.max_result = 0
+        self.past_results.clear()
+        self.progress.update(self.task_epoch, advance=-0.5)
+
+    @override
+    def on_train_batch_end(self, trainer, pl_module, outputs, batch: Any, batch_idx: int):
+        self._update(self.train_progress_bar_id, batch_idx + 1)
+        self._update_metrics(trainer, pl_module)
+        epoch_descript = "[cyan]Train [white]|"
+        batch_descript = "[green]Train [white]|"
+        metrics = self.get_metrics(trainer, pl_module)
+        metrics.pop("v_num")
+        for metrics_name, metrics_val in metrics.items():
+            if "Loss_step" in metrics_name:
+                epoch_descript += f"{metrics_name.removesuffix('_step'): ^9}|"
+                batch_descript += f"   {metrics_val:2.2f}  |"
+
+        self.progress.update(self.task_epoch, advance=1 / self.total_train_batches, description=epoch_descript)
+        self.progress.update(self.train_progress_bar_id, description=batch_descript)
         self.refresh()
-        self.remove_task(self.batch_task)
 
-    @rank_check
-    def finish_train(self):
-        self.remove_task(self.task_epoch)
-        self.stop()
-        if self.use_wandb:
-            self.wandb.finish()
-        if self.use_tensorboard:
-            self.tb_writer.close()
+    @override
+    def on_train_end(self, trainer: "Trainer", pl_module: "LightningModule") -> None:
+        self._update_metrics(trainer, pl_module)
+        self.progress.remove_task(self.train_progress_bar_id)
+        self.train_progress_bar_id = None
+
+    @override
+    def on_validation_end(self, trainer: "Trainer", pl_module: "LightningModule") -> None:
+        if trainer.state.fn == "fit":
+            self._update_metrics(trainer, pl_module)
+        self.reset_dataloader_idx_tracker()
+        all_metrics = self.get_metrics(trainer, pl_module)
+
+        ap_ar_list = [
+            key
+            for key in all_metrics.keys()
+            if key.startswith(("map", "mar")) and not key.endswith(("_step", "_epoch"))
+        ]
+        score = np.array([all_metrics[key] for key in ap_ar_list]) * 100
+
+        self.progress.table, ap_main = make_ap_table(score, self.past_results, self.max_result, trainer.current_epoch)
+        self.max_result = np.maximum(score, self.max_result)
+        self.past_results.append((trainer.current_epoch, ap_main))
+
+    @override
+    def refresh(self) -> None:
+        if self.progress:
+            self.progress.refresh()
+
+    @property
+    def validation_description(self) -> str:
+        return "[green]Validation"
+
+
+class YOLORichModelSummary(RichModelSummary):
+
+    from typing_extensions import override
+
+    @staticmethod
+    @override
+    def summarize(
+        summary_data: List[Tuple[str, List[str]]],
+        total_parameters: int,
+        trainable_parameters: int,
+        model_size: float,
+        total_training_modes: Dict[str, int],
+        **summarize_kwargs: Any,
+    ) -> None:
+        from lightning.pytorch.utilities.model_summary import get_human_readable_count
+        from rich import get_console
+        from rich.table import Table
+
+        console = get_console()
+
+        header_style: str = summarize_kwargs.get("header_style", "bold magenta")
+        table = Table(header_style=header_style)
+        table.add_column(" ", style="dim")
+        table.add_column("Name", justify="left", no_wrap=True)
+        table.add_column("Type")
+        table.add_column("Params", justify="right")
+        table.add_column("Mode")
+
+        column_names = list(zip(*summary_data))[0]
+
+        for column_name in ["In sizes", "Out sizes"]:
+            if column_name in column_names:
+                table.add_column(column_name, justify="right", style="white")
+
+        rows = list(zip(*(arr[1] for arr in summary_data)))
+        for row in rows:
+            table.add_row(*row)
+
+        console.print(table)
+
+        parameters = []
+        for param in [trainable_parameters, total_parameters - trainable_parameters, total_parameters, model_size]:
+            parameters.append("{:<{}}".format(get_human_readable_count(int(param)), 10))
+
+        grid = Table(header_style=header_style)
+        table.add_column(" ", style="dim")
+        grid.add_column("[bold]Attributes[/]")
+        grid.add_column("Value")
+
+        grid.add_row("[bold]Trainable params[/]", f"{parameters[0]}")
+        grid.add_row("[bold]Non-trainable params[/]", f"{parameters[1]}")
+        grid.add_row("[bold]Total params[/]", f"{parameters[2]}")
+        grid.add_row("[bold]Total estimated model params size (MB)[/]", f"{parameters[3]}")
+        grid.add_row("[bold]Modules in train mode[/]", f"{total_training_modes['train']}")
+        grid.add_row("[bold]Modules in eval mode[/]", f"{total_training_modes['eval']}")
+
+        console.print(grid)
+
+
+class ImageLogger(Callback):
+    def on_validation_batch_end(self, trainer: Trainer, pl_module, outputs, batch, batch_idx) -> None:
+        if batch_idx != 0:
+            return
+        batch_size, images, targets, rev_tensor, img_paths = batch
+        gt_boxes = targets[0] if targets.ndim == 3 else targets
+        pred_boxes = outputs[0] if isinstance(outputs, list) else outputs
+        images = [images[0]]
+        step = trainer.current_epoch
+        for logger in trainer.loggers:
+            if isinstance(logger, WandbLogger):
+                logger.log_image("Input Image", images, step=step)
+                logger.log_image("Ground Truth", images, step=step, boxes=[log_bbox(gt_boxes)])
+                logger.log_image("Prediction", images, step=step, boxes=[log_bbox(pred_boxes)])
+
+
+def setup(cfg: Config):
+    if hasattr(cfg, "quite"):
+        logger.removeHandler("YOLO_logger")
+        return
+
+    class EmojiFormatter(logging.Formatter):
+        def format(self, record):
+            return f":high_voltage: {super().format(record)}"
 
+    rich_handler = RichHandler(markup=True)
+    rich_handler.setFormatter(EmojiFormatter("%(message)s"))
+    lightning_logger = logging.getLogger("lightning.pytorch")
+    lightning_logger.handlers.clear()
+    lightning_logger.addHandler(rich_handler)
 
-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.info("🌐 " + line)
+    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.info(Text.from_ansi(":globe_with_meridians: " + line))
+
+    wandb.errors.term._log = custom_wandb_log
+
+    save_path = validate_log_directory(cfg, cfg.name)
+
+    progress, loggers = [], []
+    progress.append(YOLORichProgressBar())
+    progress.append(YOLORichModelSummary())
+    progress.append(ImageLogger())
+
+    loggers.append(WandbLogger(project="YOLO", name=cfg.name, save_dir=save_path, id=None))
+
+    return progress, loggers
 
 
 def log_model_structure(model: Union[ModuleList, YOLOLayer, YOLO]):
@@ -291,7 +301,7 @@ def validate_log_directory(cfg: Config, exp_name: str) -> Path:
             )
 
     save_path.mkdir(parents=True, exist_ok=True)
-    logger.info(f"📄 Created log folder: [bold gray]{save_path}[/]", extra={"markup": True})
+    logger.info(f"📄 Created log folder: [blue b u]123{save_path}[/]")
     logger.addHandler(FileHandler(save_path / "output.log"))
     return save_path
 
@@ -327,4 +337,4 @@ def log_bbox(
             bbox_entry["scores"] = {"confidence": conf[0]}
         bbox_list.append(bbox_entry)
 
-    return bbox_list
+    return {"predictions": {"box_data": bbox_list}}

yolo/utils/model_utils.py

@@ -56,23 +56,8 @@ def create_optimizer(model: YOLO, optim_cfg: OptimizerConfig) -> Optimizer:
         {"params": conv_params},
         {"params": norm_params, "weight_decay": 0},
     ]
-
-    def next_epoch(self, batch_num):
-        self.min_lr = self.max_lr
-        self.max_lr = [param["lr"] for param in self.param_groups]
-        self.batch_num = batch_num
-        self.batch_idx = 0
-
-    def next_batch(self):
-        self.batch_idx += 1
-        for lr_idx, param_group in enumerate(self.param_groups):
-            min_lr, max_lr = self.min_lr[lr_idx], self.max_lr[lr_idx]
-            param_group["lr"] = min_lr + (self.batch_idx) * (max_lr - min_lr) / self.batch_num
-
-    optimizer_class.next_batch = next_batch
-    optimizer_class.next_epoch = next_epoch
     optimizer = optimizer_class(model_parameters, **optim_cfg.args)
-    optimizer.max_lr = [0.1, 0, 0]
+    # TODO: implement batch lr schedular when warm up
     return optimizer
 
 
@@ -168,6 +153,7 @@ def predicts_to_json(img_paths, predicts, rev_tensor):
     batch_json = []
     for img_path, bboxes, box_reverse in zip(img_paths, predicts, rev_tensor):
         scale, shift = box_reverse.split([1, 4])
+        bboxes = bboxes.clone()
         bboxes[:, 1:5] = (bboxes[:, 1:5] - shift[None]) / scale[None]
         bboxes[:, 1:5] = transform_bbox(bboxes[:, 1:5], "xyxy -> xywh")
         for cls, *pos, conf in bboxes:

yolo/utils/solver_utils.py

@@ -1,5 +1,6 @@
 import contextlib
 import io
+from typing import Dict
 
 import numpy as np
 from pycocotools.coco import COCO
@@ -17,7 +18,7 @@ def calculate_ap(coco_gt: COCO, pd_path):
     return coco_eval.stats
 
 
-def make_ap_table(score, past_result=[], last_score=None, epoch=-1):
+def make_ap_table(score: Dict[str, float], past_result=[], max_result=None, epoch=-1):
     ap_table = Table()
     ap_table.add_column("Epoch", justify="center", style="white", width=5)
     ap_table.add_column("Avg. Precision", justify="left", style="cyan")
@@ -30,7 +31,7 @@ def make_ap_table(score, past_result=[], last_score=None, epoch=-1):
     if past_result:
         ap_table.add_row()
 
-    color = np.where(last_score <= score, "[green]", "[red]")
+    color = np.where(max_result <= score, "[green]", "[red]")
 
     this_ap = ("AP @ .5:.95", color[0], score[0], "AP @        .5", color[1], score[1])
     metrics = [