✨ [Add] new EMA! for torch lightning

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(
@@ -45,7 +46,7 @@ class ValidateModel(BaseModel):
 
     def validation_step(self, batch, batch_idx):
         batch_size, images, targets, rev_tensor, img_paths = batch
-        predicts = self.post_process(self(images), image_size=images.shape[2:])
+        predicts = self.post_process(self.ema(images), image_size=images.shape[2:])
         batch_metrics = self.metric(
             [to_metrics_format(predict) for predict in predicts], [to_metrics_format(target) for target in targets]
         )

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
 
 
@@ -255,6 +256,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
 
@@ -31,28 +36,31 @@ def lerp(start: float, end: float, step: Union[int, float], total: int = 1):
     return start + (end - start) * step / total
 
 
-class ExponentialMovingAverage:
-    def __init__(self, model: torch.nn.Module, decay: float):
-        self.model = model
+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()}
-
-    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)
+        self.tau = tau
+        self.step = 0
+
+    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()]
+
+    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)
+            if dist.is_initialized():
+                dist.broadcast(param, src=0)
+
+    @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: