🔀 [Merge] branch 'TRAIN' into TEST

yolo/config/config.py

@@ -97,7 +97,7 @@ class SchedulerConfig:
 
 @dataclass
 class EMAConfig:
-    enabled: bool
+    enable: bool
     decay: float
 
 

yolo/config/task/train.yaml

@@ -50,5 +50,5 @@ scheduler:
     end_factor: 0.01
 
 ema:
-  enabled: true
+  enable: true
   decay: 0.995

yolo/config/task/validation.yaml

@@ -7,7 +7,7 @@ data:
   shuffle: False
   pin_memory: True
   data_augment: {}
-  dynamic_shape: True
+  dynamic_shape: False
 nms:
   min_confidence: 0.0001
   min_iou: 0.7

yolo/tools/data_loader.py

@@ -56,7 +56,15 @@ class YoloDataset(Dataset):
             data = self.filter_data(dataset_path, phase_name, self.dynamic_shape)
             torch.save(data, cache_path)
         else:
-            data = torch.load(cache_path, weights_only=False)
+            try:
+                data = torch.load(cache_path, weights_only=False)
+            except Exception as e:
+                logger.error(
+                    f":rotating_light: Failed to load the cache at '{cache_path}'.\n"
+                    ":rotating_light: This may be caused by using cache from different other YOLO.\n"
+                    ":rotating_light: Please clean the cache and try running again."
+                )
+                raise e
             logger.info(f":package: Loaded {phase_name} cache")
         return data
 

yolo/tools/solver.py

@@ -33,6 +33,7 @@ class ValidateModel(BaseModel):
         self.metric = MeanAveragePrecision(iou_type="bbox", box_format="xyxy")
         self.metric.warn_on_many_detections = False
         self.val_loader = create_dataloader(self.validation_cfg.data, self.cfg.dataset, self.validation_cfg.task)
+        self.ema = self.model
 
     def setup(self, stage):
         self.vec2box = create_converter(
@@ -46,7 +47,7 @@ class ValidateModel(BaseModel):
     def validation_step(self, batch, batch_idx):
         batch_size, images, targets, rev_tensor, img_paths = batch
         H, W = images.shape[2:]
-        predicts = self.post_process(self(images), image_size=[W, H])
+        predicts = self.post_process(self.ema(images), image_size=[W, H])
         batch_metrics = self.metric(
             [to_metrics_format(predict) for predict in predicts], [to_metrics_format(target) for target in targets]
         )
@@ -56,7 +57,6 @@ class ValidateModel(BaseModel):
                 "map": batch_metrics["map"],
                 "map_50": batch_metrics["map_50"],
             },
-            on_step=True,
             batch_size=batch_size,
         )
         return predicts
@@ -64,9 +64,11 @@ class ValidateModel(BaseModel):
     def on_validation_epoch_end(self):
         epoch_metrics = self.metric.compute()
         del epoch_metrics["classes"]
-        self.log_dict(epoch_metrics, prog_bar=True, rank_zero_only=True)
+        self.log_dict(epoch_metrics, prog_bar=True, sync_dist=True, rank_zero_only=True)
         self.log_dict(
-            {"PyCOCO/AP @ .5:.95": epoch_metrics["map"], "PyCOCO/AP @ .5": epoch_metrics["map_50"]}, rank_zero_only=True
+            {"PyCOCO/AP @ .5:.95": epoch_metrics["map"], "PyCOCO/AP @ .5": epoch_metrics["map_50"]},
+            sync_dist=True,
+            rank_zero_only=True,
         )
         self.metric.reset()
 
@@ -85,7 +87,9 @@ class TrainModel(ValidateModel):
         return self.train_loader
 
     def on_train_epoch_start(self):
-        self.trainer.optimizers[0].next_epoch(ceil(len(self.train_loader) / self.trainer.world_size))
+        self.trainer.optimizers[0].next_epoch(
+            ceil(len(self.train_loader) / self.trainer.world_size), self.current_epoch
+        )
         self.vec2box.update(self.cfg.image_size)
 
     def training_step(self, batch, batch_idx):

yolo/utils/bounding_box_utils.py

@@ -212,19 +212,20 @@ class BoxMatcher:
         topk_masks = topk_targets > 0
         return topk_targets, topk_masks
 
-    def filter_duplicates(self, target_matrix: Tensor, topk_mask: Tensor):
+    def filter_duplicates(self, iou_mat: Tensor, topk_mask: Tensor, grid_mask: Tensor):
         """
         Filter the maximum suitability target index of each anchor.
 
         Args:
-            target_matrix [batch x targets x anchors]: The suitability for each targets-anchors
+            iou_mat [batch x targets x anchors]: The suitability for each targets-anchors
 
         Returns:
             unique_indices [batch x anchors x 1]: The index of the best targets for each anchors
         """
         duplicates = (topk_mask.sum(1, keepdim=True) > 1).repeat([1, topk_mask.size(1), 1])
-        max_idx = F.one_hot(target_matrix.argmax(1), topk_mask.size(1)).permute(0, 2, 1)
+        max_idx = F.one_hot(iou_mat.argmax(1), topk_mask.size(1)).permute(0, 2, 1)
         topk_mask = torch.where(duplicates, max_idx, topk_mask)
+        topk_mask &= grid_mask
         unique_indices = topk_mask.argmax(dim=1)
         return unique_indices[..., None], topk_mask.sum(1), topk_mask
 
@@ -278,7 +279,7 @@ class BoxMatcher:
         topk_targets, topk_mask = self.filter_topk(target_matrix, topk=self.topk)
 
         # delete one anchor pred assign to mutliple gts
-        unique_indices, valid_mask, topk_mask = self.filter_duplicates(iou_mat, topk_mask)
+        unique_indices, valid_mask, topk_mask = self.filter_duplicates(iou_mat, topk_mask, grid_mask)
 
         align_bbox = torch.gather(target_bbox, 1, unique_indices.repeat(1, 1, 4))
         align_cls = torch.gather(target_cls, 1, unique_indices).squeeze(-1)

yolo/utils/dataset_utils.py

@@ -115,7 +115,14 @@ def scale_segmentation(
 
 
 def tensorlize(data):
-    img_paths, bboxes, img_ratios = zip(*data)
+    try:
+        img_paths, bboxes, img_ratios = zip(*data)
+    except ValueError as e:
+        logger.error(
+            ":rotating_light: This may be caused by using old cache or another version of YOLO's cache.\n"
+            ":rotating_light: Please clean the cache and try running again."
+        )
+        raise e
     max_box = max(bbox.size(0) for bbox in bboxes)
     padded_bbox_list = []
     for bbox in bboxes:

yolo/utils/logging_utils.py

@@ -38,6 +38,7 @@ from typing_extensions import override
 from yolo.config.config import Config, YOLOLayer
 from yolo.model.yolo import YOLO
 from yolo.utils.logger import logger
+from yolo.utils.model_utils import EMA
 from yolo.utils.solver_utils import make_ap_table
 
 
@@ -97,7 +98,6 @@ class YOLORichProgressBar(RichProgressBar):
         )
         self.max_result = 0
         self.past_results.clear()
-        self.progress.update(self.task_epoch, advance=-0.5)
 
     @override
     @rank_zero_only
@@ -255,6 +255,8 @@ def setup(cfg: Config):
 
     progress, loggers = [], []
 
+    if hasattr(cfg.task, "ema") and cfg.task.ema.enable:
+        progress.append(EMA(cfg.task.ema.decay))
     if quite:
         logger.setLevel(logging.ERROR)
         return progress, loggers, save_path

yolo/utils/model_utils.py

@@ -1,11 +1,16 @@
 import os
+from copy import deepcopy
+from math import exp
 from pathlib import Path
 from typing import List, Optional, Type, Union
 
 import torch
 import torch.distributed as dist
+from lightning import LightningModule, Trainer
+from lightning.pytorch.callbacks import Callback
+from lightning.pytorch.utilities import rank_zero_only
 from omegaconf import ListConfig
-from torch import Tensor
+from torch import Tensor, no_grad
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LambdaLR, SequentialLR, _LRScheduler
 
@@ -15,28 +20,48 @@ from yolo.utils.bounding_box_utils import Anc2Box, Vec2Box, bbox_nms, transform_
 from yolo.utils.logger import logger
 
 
-class ExponentialMovingAverage:
-    def __init__(self, model: torch.nn.Module, decay: float):
-        self.model = model
+def lerp(start: float, end: float, step: Union[int, float], total: int = 1):
+    """
+    Linearly interpolates between start and end values.
+
+    Parameters:
+        start (float): The starting value.
+        end (float): The ending value.
+        step (int): The current step in the interpolation process.
+        total (int): The total number of steps.
+
+    Returns:
+        float: The interpolated value.
+    """
+    return start + (end - start) * step / total
+
+
+class EMA(Callback):
+    def __init__(self, decay: float = 0.9999, tau: float = 500):
+        super().__init__()
+        logger.info(":chart_with_upwards_trend: Enable Model EMA")
         self.decay = decay
-        self.shadow = {name: param.clone().detach() for name, param in model.named_parameters()}
+        self.tau = tau
+        self.step = 0
 
-    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 setup(self, trainer, pl_module, stage):
+        pl_module.ema = deepcopy(pl_module.model)
+        self.ema_parameters = [param.clone().detach().to(pl_module.device) for param in pl_module.parameters()]
+        for param in pl_module.ema.parameters():
+            param.requires_grad = False
 
-    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 on_validation_start(self, trainer: "Trainer", pl_module: "LightningModule"):
+        for param, ema_param in zip(pl_module.ema.parameters(), self.ema_parameters):
+            param.data.copy_(ema_param)
+            trainer.strategy.broadcast(param)
 
-    def restore(self):
-        """Restore the original parameters from the shadow."""
-        for name, param in self.model.named_parameters():
-            self.shadow[name].copy_(param.data)
+    @rank_zero_only
+    @no_grad()
+    def on_train_batch_end(self, trainer: "Trainer", pl_module: "LightningModule", *args, **kwargs) -> None:
+        self.step += 1
+        decay_factor = self.decay * (1 - exp(-self.step / self.tau))
+        for param, ema_param in zip(pl_module.parameters(), self.ema_parameters):
+            ema_param.data.copy_(lerp(param.detach(), ema_param, decay_factor))
 
 
 def create_optimizer(model: YOLO, optim_cfg: OptimizerConfig) -> Optimizer:
@@ -57,9 +82,15 @@ def create_optimizer(model: YOLO, optim_cfg: OptimizerConfig) -> Optimizer:
         {"params": norm_params, "momentum": 0.8, "weight_decay": 0},
     ]
 
-    def next_epoch(self, batch_num):
+    def next_epoch(self, batch_num, epoch_idx):
         self.min_lr = self.max_lr
         self.max_lr = [param["lr"] for param in self.param_groups]
+        # TODO: load momentum from config instead a fix number
+        #       0.937: Start Momentum
+        #       0.8  : Normal Momemtum
+        #       3    : The warm up epoch num
+        self.min_mom = lerp(0.937, 0.8, max(epoch_idx, 3), 3)
+        self.max_mom = lerp(0.937, 0.8, max(epoch_idx + 1, 3), 3)
         self.batch_num = batch_num
         self.batch_idx = 0
 
@@ -68,7 +99,8 @@ def create_optimizer(model: YOLO, optim_cfg: OptimizerConfig) -> Optimizer:
         lr_dict = dict()
         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
+            param_group["lr"] = lerp(min_lr, max_lr, self.batch_idx, self.batch_num)
+            param_group["momentum"] = lerp(self.min_mom, self.max_mom, self.batch_idx, self.batch_num)
             lr_dict[f"LR/{lr_idx}"] = param_group["lr"]
         return lr_dict