Fix: resume_training: True and EMACallback active (#372)

Co-authored-by: Martin Kozlovsky <martin.kozlovsky@luxonis.com>

Author: dtronmans

luxonis_train/callbacks/ema.py

@@ -1,12 +1,16 @@
 import math
+from collections.abc import Mapping
 from copy import deepcopy
 from typing import Any
 
 import lightning.pytorch as pl
 import torch
 from lightning.pytorch.utilities.types import STEP_OUTPUT
+from loguru import logger
 from torch import nn
 
+from luxonis_train.utils.checkpoint import filter_checkpoint_state_dict
+
 
 class ModelEma(nn.Module):
     """Model Exponential Moving Average.
@@ -65,13 +69,13 @@ def update(self, model: pl.LightningModule) -> None:
             else:
                 decay = self.decay
 
+            model_state_dict = model.state_dict()
             ema_lerp_values = []
             model_lerp_values = []
-            for ema_v, model_v in zip(
-                self.state_dict_ema.values(),
-                model.state_dict().values(),
-                strict=True,
-            ):
+            for key, ema_v in self.state_dict_ema.items():
+                model_v = model_state_dict.get(key)
+                if model_v is None:
+                    continue
                 if ema_v.is_floating_point():
                     ema_lerp_values.append(ema_v)
                     model_lerp_values.append(model_v)
@@ -115,8 +119,13 @@ def __init__(
 
         self._ema = None
         self.loaded_ema_state_dict = None
+        self.loaded_ema_updates = None
         self.collected_state_dict = None
 
+    @staticmethod
+    def _format_key_list(keys: set[str]) -> str:
+        return ", ".join(sorted(keys)) if keys else "<none>"
+
     @property
     def ema(self) -> ModelEma:
         if self._ema is None:
@@ -144,12 +153,54 @@ def on_fit_start(
             target_device = next(
                 iter(self._ema.state_dict_ema.values())
             ).device
-            self.loaded_ema_state_dict = {
-                k: v.to(target_device)
-                for k, v in self.loaded_ema_state_dict.items()
+            current_state_dict = self._ema.state_dict_ema
+            comparable_current_state_dict = filter_checkpoint_state_dict(
+                current_state_dict
+            )
+            comparable_loaded_state_dict = filter_checkpoint_state_dict(
+                self.loaded_ema_state_dict
+            )
+            current_keys = set(comparable_current_state_dict)
+            loaded_keys = set(comparable_loaded_state_dict)
+            missing_in_checkpoint = current_keys - loaded_keys
+            extra_in_checkpoint = loaded_keys - current_keys
+            incompatible_shapes = {
+                key
+                for key in current_keys & loaded_keys
+                if comparable_current_state_dict[key].shape
+                != comparable_loaded_state_dict[key].shape
             }
-            self._ema.state_dict_ema = self.loaded_ema_state_dict
+
+            if missing_in_checkpoint:
+                logger.warning(
+                    "EMA checkpoint is missing keys present in the current model. "
+                    "Keeping freshly initialized EMA values for: "
+                    f"{self._format_key_list(missing_in_checkpoint)}"
+                )
+            if extra_in_checkpoint:
+                logger.warning(
+                    "EMA checkpoint contains keys not present in the current model. "
+                    "Ignoring: "
+                    f"{self._format_key_list(extra_in_checkpoint)}"
+                )
+            if incompatible_shapes:
+                logger.warning(
+                    "EMA checkpoint contains keys with incompatible shapes. "
+                    "Ignoring: "
+                    f"{self._format_key_list(incompatible_shapes)}"
+                )
+
+            for key, value in comparable_loaded_state_dict.items():
+                if (
+                    key in current_state_dict
+                    and key not in incompatible_shapes
+                ):
+                    current_state_dict[key] = value.to(target_device)
+            self._ema.state_dict_ema = current_state_dict
+            if self.loaded_ema_updates is not None:
+                self._ema.updates = self.loaded_ema_updates
             self.loaded_ema_state_dict = None
+            self.loaded_ema_updates = None
 
     def on_train_batch_end(
         self,
@@ -248,7 +299,7 @@ def on_save_checkpoint(
         trainer: pl.Trainer,
         pl_module: pl.LightningModule,
         checkpoint: dict,
-    ) -> None:  # or dict?
+    ) -> None:
         """Save the EMA state dictionary into the checkpoint.
 
         @type trainer: L{pl.Trainer}
@@ -261,6 +312,19 @@ def on_save_checkpoint(
         if self._ema is not None:
             checkpoint["state_dict"] = self._ema.state_dict_ema
 
+    def state_dict(self) -> dict[str, Any]:
+        if self._ema is None:
+            return {}
+        return {
+            "ema_state_dict": filter_checkpoint_state_dict(
+                self._ema.state_dict_ema
+            ),
+            "updates": self._ema.updates,
+        }
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        self._load_ema_state(state_dict)
+
     def on_load_checkpoint(
         self,
         trainer: pl.Trainer,
@@ -272,8 +336,18 @@ def on_load_checkpoint(
         @type callback_state: dict
         @param callback_state: Pytorch Lightning callback state.
         """
-        if callback_state and "state_dict" in callback_state:
-            self.loaded_ema_state_dict = callback_state["state_dict"]
+        self._load_ema_state(callback_state)
+
+    def _load_ema_state(self, state_dict: dict[str, Any]) -> None:
+        if state_dict:
+            loaded_state_dict = state_dict.get(
+                "ema_state_dict", state_dict.get("state_dict")
+            )
+            if isinstance(loaded_state_dict, Mapping):
+                self.loaded_ema_state_dict = loaded_state_dict
+            updates = state_dict.get("updates")
+            if isinstance(updates, int):
+                self.loaded_ema_updates = updates
 
     def _swap_to_ema_weights(self, pl_module: pl.LightningModule) -> None:
         """Swap the current model weights with the EMA weights.

luxonis_train/lightning/luxonis_lightning.py

@@ -1,4 +1,3 @@
-import re
 from collections import defaultdict
 from collections.abc import Callable, Mapping
 from pathlib import Path
@@ -26,6 +25,7 @@
 from luxonis_train.nodes import BaseNode
 from luxonis_train.typing import Labels, Packet
 from luxonis_train.utils import DatasetMetadata, LuxonisTrackerPL
+from luxonis_train.utils.checkpoint import filter_checkpoint_state_dict
 
 from .luxonis_output import LuxonisOutput
 from .utils import (
@@ -1027,14 +1027,9 @@ def _strip_state_prefix(key: str) -> str:
     def _add_custom_data_to_checkpoint(
         self, checkpoint: dict[str, Any]
     ) -> None:
-        pattern = re.compile(
-            r"^nodes\.[^.]+\.(metrics|visualizers|losses)\..*_node\..*"
+        checkpoint["state_dict"] = filter_checkpoint_state_dict(
+            checkpoint["state_dict"]
         )
-        checkpoint["state_dict"] = {
-            k: v
-            for k, v in checkpoint["state_dict"].items()
-            if not pattern.match(k)
-        }
         checkpoint |= {
             "version": luxonis_train.__version__,
             "execution_order": get_model_execution_order(self),

luxonis_train/lightning/utils.py

@@ -493,16 +493,6 @@ def build_callbacks(
                 "in the callbacks list. The `accumulate_grad_batches` "
                 "parameter in the config will be ignored."
             )
-    callbacks.append(
-        ModelCheckpoint(
-            dirpath=save_dir / "min_val_loss",
-            filename=f"{model_name}_loss={{val/loss:.4f}}_{{epoch:02d}}",
-            monitor="val/loss",
-            auto_insert_metric_name=False,
-            save_top_k=cfg.trainer.save_top_k,
-            mode="min",
-        ),
-    )
     if main_metric is not None:
         node_name, metric_name = main_metric
         formatted_node = nodes.formatted_name(node_name)
@@ -521,6 +511,16 @@ def build_callbacks(
             )
         )
 
+    callbacks.append(
+        ModelCheckpoint(
+            dirpath=save_dir / "min_val_loss",
+            filename=f"{model_name}_loss={{val/loss:.4f}}_{{epoch:02d}}",
+            monitor="val/loss",
+            auto_insert_metric_name=False,
+            save_top_k=cfg.trainer.save_top_k,
+            mode="min",
+        )
+    )
     return callbacks
 
 

luxonis_train/utils/__init__.py

@@ -9,6 +9,10 @@
     keypoints_to_bboxes,
     non_max_suppression,
 )
+from .checkpoint import (
+    CHECKPOINT_FILTERED_STATE_DICT_PATTERN,
+    filter_checkpoint_state_dict,
+)
 from .dataset_metadata import DatasetMetadata
 from .exceptions import IncompatibleError
 from .general import (
@@ -41,6 +45,7 @@
 from .tracker import LuxonisTrackerPL
 
 __all__ = [
+    "CHECKPOINT_FILTERED_STATE_DICT_PATTERN",
     "Counter",
     "DatasetMetadata",
     "IncompatibleError",
@@ -55,6 +60,7 @@
     "compute_pose_oks",
     "default_annotate",
     "dist2bbox",
+    "filter_checkpoint_state_dict",
     "get_attribute_check_none",
     "get_batch_instances",
     "get_batch_instances",

luxonis_train/utils/checkpoint.py

@@ -0,0 +1,18 @@
+import re
+from collections.abc import Mapping
+
+from torch import Tensor
+
+CHECKPOINT_FILTERED_STATE_DICT_PATTERN = re.compile(
+    r"^nodes\.[^.]+\.(metrics|visualizers|losses)\..*_node\..*"
+)
+
+
+def filter_checkpoint_state_dict(
+    state_dict: Mapping[str, Tensor],
+) -> dict[str, Tensor]:
+    return {
+        key: value
+        for key, value in state_dict.items()
+        if not CHECKPOINT_FILTERED_STATE_DICT_PATTERN.match(key)
+    }

tests/integration/test_resume_training.py

@@ -0,0 +1,52 @@
+from pathlib import Path
+
+from luxonis_ml.data import LuxonisDataset
+from luxonis_ml.typing import Params
+
+from luxonis_train.core import LuxonisModel
+
+
+def test_resume_training_with_ema_does_not_crash(
+    parking_lot_dataset: LuxonisDataset, opts: Params, tmp_path: Path
+):
+    config_file = "configs/detection_light_model.yaml"
+    save_dir = tmp_path / "save-directory"
+
+    train_opts = opts | {
+        "loader.params.dataset_name": parking_lot_dataset.identifier,
+        "loader.train_view": "train",
+        "loader.val_view": "train",
+        "loader.test_view": "train",
+        "model.predefined_model.params.task_name": "vehicles",
+        "trainer.overfit_batches": 1,
+        "trainer.seed": 42,
+        "trainer.deterministic": "warn",
+        "trainer.epochs": 1,
+        "trainer.validation_interval": 1,
+        "tracker.save_directory": str(save_dir),
+        "trainer.callbacks": [
+            {
+                "name": "EMACallback",
+                "active": True,
+                "params": {"decay": 0.9999},
+            },
+            {"name": "TestOnTrainEnd", "active": False},
+            {"name": "ExportOnTrainEnd", "active": False},
+            {"name": "ArchiveOnTrainEnd", "active": False},
+            {"name": "ConvertOnTrainEnd", "active": False},
+            {"name": "UploadCheckpoint", "active": False},
+        ],
+    }
+
+    model = LuxonisModel(config_file, train_opts)
+    model.train()
+
+    ckpt_path = model.get_best_metric_checkpoint_path()
+    assert ckpt_path, "No checkpoint found after initial training"
+
+    resume_opts = train_opts | {
+        "trainer.resume_training": True,
+        "trainer.epochs": 2,
+    }
+    resumed_model = LuxonisModel(config_file, resume_opts)
+    resumed_model.train(weights=ckpt_path)

tests/unittests/test_callbacks/test_ema.py

@@ -97,6 +97,10 @@ def test_ema_state_saved_to_checkpoint(
     ema_callback.on_save_checkpoint(trainer, model, checkpoint)
 
     assert "state_dict" in checkpoint
+    assert (
+        checkpoint["state_dict"].keys()
+        == ema_callback.ema.state_dict_ema.keys()
+    )
 
 
 def test_load_from_checkpoint(