Add EMA callback (#74)

* 👽 pytorch WeightAveraging callback which will become available in pytorch-lightning 2.5.4

* 🚨 Fix linting issues in the WeightAveraging callback.

* ✨ Add an EMAWeightAveragingCallback that uses the upstream WeightAveraging callback but has a simpler configuration interface

* 🔧 Add the EMAWeightAveragingCallback to our set of default callbacks

* ✨ Make decay rate configurable

Author: jemrobinson

ice_station_zebra/callbacks/__init__.py

@@ -1,5 +1,11 @@
+from .ema_weight_averaging_callback import EMAWeightAveragingCallback
 from .metric_summary_callback import MetricSummaryCallback
 from .plotting_callback import PlottingCallback
 from .unconditional_checkpoint import UnconditionalCheckpoint
 
-__all__ = ["MetricSummaryCallback", "PlottingCallback", "UnconditionalCheckpoint"]
+__all__ = [
+    "EMAWeightAveragingCallback",
+    "MetricSummaryCallback",
+    "PlottingCallback",
+    "UnconditionalCheckpoint",
+]

ice_station_zebra/callbacks/ema_weight_averaging_callback.py

@@ -0,0 +1,40 @@
+from torch.optim.swa_utils import get_ema_multi_avg_fn
+
+from .weight_averaging import WeightAveraging
+
+
+class EMAWeightAveragingCallback(WeightAveraging):
+    """A callback that updates an averaged model for Exponential Moving Average (EMA) after each training step."""
+
+    def __init__(
+        self,
+        *,
+        decay_rate: float,
+        every_n_epochs: int | None = None,
+        every_n_steps: int | None = None,
+    ) -> None:
+        """Summarise metrics during evaluation.
+
+        Args:
+            decay_rate: Parameter update decay rate.
+            every_n_epochs: How many epochs to wait before updating.
+            every_n_steps: How many steps to wait before updating.
+
+        """
+        super().__init__(
+            multi_avg_fn=get_ema_multi_avg_fn(decay_rate), use_buffers=True
+        )
+        self.every_n_epochs = every_n_epochs
+        self.every_n_steps = every_n_steps
+
+    def should_update(
+        self, step_idx: int | None = None, epoch_idx: int | None = None
+    ) -> bool:
+        """Update if we are at the requested number of steps or epochs."""
+        if self.every_n_epochs and epoch_idx:
+            return epoch_idx % self.every_n_epochs == 0
+
+        if self.every_n_steps and step_idx:
+            return step_idx % self.every_n_steps == 0
+
+        return False