| from typing import Any, Optional, Union |
|
|
| import torch |
| from torch.optim.swa_utils import get_ema_avg_fn |
| from lightning.pytorch.utilities.rank_zero import rank_zero_info |
|
|
| from src.callbacks.lightning_weight_averaging import WeightAveraging |
|
|
|
|
| class WarmupEMAWeightAveraging(WeightAveraging): |
| def __init__( |
| self, |
| warmup_pct: float, |
| enabled: bool = True, |
| decay: Optional[float] = None, |
| decay_numerator: float = 20.0, |
| update_every_n_steps: int = 1, |
| update_starting_at_step: Optional[int] = None, |
| device: Optional[Union[torch.device, str, int]] = None, |
| use_buffers: bool = True, |
| ) -> None: |
| super().__init__(device=device, use_buffers=use_buffers) |
| self.enabled = enabled |
| self.warmup_pct = warmup_pct |
| self.decay = decay |
| self.decay_numerator = decay_numerator |
| self.update_every_n_steps = update_every_n_steps |
| self.update_starting_at_step = update_starting_at_step |
|
|
| self.resolved_decay: Optional[float] = None |
| self.resolved_start_step: Optional[int] = None |
|
|
| def setup(self, trainer, pl_module, stage: str) -> None: |
| if not self.enabled: |
| rank_zero_info("WarmupEMAWeightAveraging is disabled by config.") |
| return |
|
|
| if stage != "fit": |
| return super().setup(trainer, pl_module, stage) |
|
|
| if trainer.max_steps and trainer.max_steps > 0: |
| total_steps = trainer.max_steps |
| else: |
| total_steps = trainer.estimated_stepping_batches |
|
|
| if total_steps <= 0: |
| total_steps = 100000 |
|
|
| warmup_steps = int(total_steps * self.warmup_pct) |
| if self.update_starting_at_step is None: |
| self.resolved_start_step = warmup_steps |
| else: |
| self.resolved_start_step = self.update_starting_at_step |
|
|
| if self.decay is None: |
| active_steps = max(1, total_steps - self.resolved_start_step) |
| computed_decay = 1.0 - (self.decay_numerator / active_steps) |
| self.resolved_decay = min(0.99999, max(0.9, computed_decay)) |
| else: |
| self.resolved_decay = self.decay |
|
|
| self._kwargs["avg_fn"] = get_ema_avg_fn(decay=self.resolved_decay) |
|
|
| super().setup(trainer, pl_module, stage) |
|
|
| rank_zero_info( |
| "WarmupEMAWeightAveraging configured: " |
| f"total_steps={total_steps}, warmup_steps={warmup_steps}, " |
| f"start_step={self.resolved_start_step}, decay={self.resolved_decay:.8f}" |
| ) |
|
|
| def should_update( |
| self, step_idx: Optional[int] = None, epoch_idx: Optional[int] = None |
| ) -> bool: |
| if step_idx is None: |
| return False |
|
|
| if not self.enabled: |
| return False |
|
|
| if self.resolved_start_step is None: |
| return False |
|
|
| if step_idx < self.resolved_start_step: |
| return False |
|
|
| if self.update_every_n_steps <= 0: |
| return False |
|
|
| return step_idx % self.update_every_n_steps == 0 |
|
|
| def state_dict(self) -> dict[str, Any]: |
| state = super().state_dict() |
| state["resolved_decay"] = self.resolved_decay |
| state["resolved_start_step"] = self.resolved_start_step |
| return state |
|
|
| def load_state_dict(self, state_dict: dict[str, Any]) -> None: |
| super().load_state_dict(state_dict) |
| self.resolved_decay = state_dict.get("resolved_decay") |
| self.resolved_start_step = state_dict.get("resolved_start_step") |
|
|
