Fix loading/saving checkpoints in OTX (#4433)

* Minor fix to torch.load() (#4392)

* fix weights only

* change otx version

* change Changelog

* fix linter

* fix checkpoint loading

* fix resume

* fix exort/testing

* try to import error from pickle separately

* address bandit issue

* add test to check backward compatability

* continue add test

* add tests

* add snapshots

* fix unit tests

* try full path

* try absolute

* updated snapshots

* bare except

* quick test

* quick test to see an error

* finally fix the error

* fix integration test

* fix integratiom tests for SS, Anomaly and CLS

* fix instance seg

Author: kprokofi

src/otx/core/metrics/fmeasure.py

@@ -761,6 +761,10 @@ def f_measure_per_confidence(self) -> None | dict:
     @property
     def best_confidence_threshold(self) -> float:
         """Returns best confidence threshold as ScoreMetric if exists."""
+        if isinstance(self._best_confidence_threshold, np.float_):
+            # Convert numpy float to python float
+            self._best_confidence_threshold = self._best_confidence_threshold.item()
+
         if self._best_confidence_threshold is None:
             msg = (
                 "Cannot obtain best_confidence_threshold updated previously. "

src/otx/core/model/anomaly.py

@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+from dataclasses import asdict
 from typing import TYPE_CHECKING, Any, Sequence, TypeAlias
 
 import torch
@@ -344,6 +345,10 @@ def get_dummy_input(self, batch_size: int = 1) -> AnomalyModelInputs:
         msg = "Wrong anomaly task type"
         raise RuntimeError(msg)
 
+    @staticmethod
+    def _dispatch_label_info(*args, **kwargs) -> AnomalyLabelInfo:  # noqa: ARG004
+        return AnomalyLabelInfo()
+
 
 class AnomalyMixin:
     """Mixin inherited before AnomalibModule to override OTXModel methods."""
@@ -437,9 +442,9 @@ def on_save_checkpoint(self, checkpoint: dict[str, Any]) -> None:
         # calls Anomalib's on_save_checkpoint
         super().on_save_checkpoint(checkpoint)  # type: ignore[misc]
 
-        checkpoint["label_info"] = self.label_info  # type: ignore[attr-defined]
+        checkpoint["hyper_parameters"]["label_info"] = asdict(self.label_info)  # type: ignore[attr-defined]
         checkpoint["otx_version"] = __version__
-        checkpoint["tile_config"] = self.tile_config  # type: ignore[attr-defined]
+        checkpoint["hyper_parameters"]["tile_config"] = asdict(self.tile_config)  # type: ignore[attr-defined]
 
         attrs = ["_input_shape", "image_threshold", "pixel_threshold"]
         checkpoint["anomaly"] = {key: getattr(self, key, None) for key in attrs}
@@ -451,3 +456,21 @@ def on_load_checkpoint(self, checkpoint: dict[str, Any]) -> None:
         if anomaly_attrs := checkpoint.get("anomaly"):
             for key, value in anomaly_attrs.items():
                 setattr(self, key, value)
+
+    def load_state_dict(self, ckpt: dict[str, Any], *args, **kwargs) -> None:
+        """Load state dictionary from checkpoint state dictionary.
+
+        If checkpoint's label_info and OTXLitModule's label_info are different,
+        load_state_pre_hook for smart weight loading will be registered.
+        """
+        self.on_load_checkpoint(ckpt)  # type: ignore[misc]
+        state_dict = ckpt["state_dict"]
+        if "label_info" in ckpt.get("hyper_parameters", {}):
+            label_info = self._dispatch_label_info(ckpt["hyper_parameters"]["label_info"])  # type: ignore[attr-defined]
+            if label_info != self.label_info:  # type: ignore[attr-defined]
+                msg = f"Checkpoint's label_info {label_info} "
+                "does not match with OTXLitModule's label_info {self.label_info}"  # type: ignore[attr-defined]
+                raise ValueError(
+                    msg,
+                )
+        return super().load_state_dict(state_dict, *args, **kwargs)  # type: ignore[misc]

src/otx/core/model/base.py

@@ -12,6 +12,7 @@
 import logging
 import warnings
 from abc import abstractmethod
+from dataclasses import asdict
 from typing import TYPE_CHECKING, Any, Callable, Literal, Sequence
 
 import numpy as np
@@ -106,13 +107,13 @@ class OTXModel(LightningModule):
 
     def __init__(
         self,
-        label_info: LabelInfoTypes,
+        label_info: LabelInfoTypes | dict,
         input_size: tuple[int, int] | None = None,
         optimizer: OptimizerCallable = DefaultOptimizerCallable,
         scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable,
         metric: MetricCallable = NullMetricCallable,
         torch_compile: bool = False,
-        tile_config: TileConfig = TileConfig(enable_tiler=False),
+        tile_config: TileConfig | dict = TileConfig(enable_tiler=False),
     ) -> None:
         super().__init__()
 
@@ -121,21 +122,21 @@ def __init__(
         self.input_size = input_size
         self.classification_layers: dict[str, dict[str, Any]] = {}
         self.model = self._create_model()
-        self.optimizer_callable = ensure_callable(optimizer)
-        self.scheduler_callable = ensure_callable(scheduler)
-        self.metric_callable = ensure_callable(metric)
+        self.optimizer_callable: OptimizerCallable = ensure_callable(optimizer)
+        self.scheduler_callable: LRSchedulerCallable = ensure_callable(scheduler)
+        self.metric_callable: MetricCallable = ensure_callable(metric)
 
         self.torch_compile = torch_compile
         self._explain_mode = False
 
         # NOTE: To guarantee immutablility of the default value
+        if isinstance(tile_config, dict):
+            tile_config = TileConfig(**tile_config)
         self._tile_config = tile_config.clone()
-
-        # this line allows to access init params with 'self.hparams' attribute
-        # also ensures init params will be stored in ckpt
-        # TODO(vinnamki): Ticket no. 138995: MetricCallable should be saved in the checkpoint
-        # so that it can retrieve it from the checkpoint
-        self.save_hyperparameters(logger=False, ignore=["optimizer", "scheduler", "metric"])
+        self.save_hyperparameters(
+            logger=False,
+            ignore=["optimizer", "scheduler", "metric", "label_info", "tile_config"],
+        )
 
     def training_step(self, batch: T_OTXBatchDataEntity, batch_idx: int) -> Tensor | None:
         """Step for model training."""
@@ -376,38 +377,42 @@ def on_save_checkpoint(self, checkpoint: dict[str, Any]) -> None:
             compiled_state_dict = checkpoint["state_dict"]
             checkpoint["state_dict"] = remove_state_dict_prefix(compiled_state_dict, "_orig_mod.")
         super().on_save_checkpoint(checkpoint)
-
-        checkpoint["label_info"] = self.label_info
+        checkpoint["hyper_parameters"]["label_info"] = asdict(self.label_info)
         checkpoint["otx_version"] = __version__
-        checkpoint["tile_config"] = self.tile_config
+        checkpoint["hyper_parameters"]["tile_config"] = asdict(self.tile_config)
+        checkpoint.pop("datamodule_hparams_name", None)
+        checkpoint.pop(
+            "datamodule_hyper_parameters",
+            None,
+        )  # Remove datamodule_hyper_parameters to prevent storing OTX classes
 
     def on_load_checkpoint(self, checkpoint: dict[str, Any]) -> None:
         """Callback on loading checkpoint."""
         super().on_load_checkpoint(checkpoint)
-
-        if ckpt_label_info := checkpoint.get("label_info"):
-            if isinstance(ckpt_label_info, LabelInfo) and not hasattr(ckpt_label_info, "label_ids"):
-                # NOTE: This is for backward compatibility
-                ckpt_label_info = LabelInfo(
-                    label_groups=ckpt_label_info.label_groups,
-                    label_names=ckpt_label_info.label_names,
-                    label_ids=ckpt_label_info.label_names,
-                )
-            self._label_info = ckpt_label_info
-
-        if ckpt_tile_config := checkpoint.get("tile_config"):
-            self.tile_config = ckpt_tile_config
+        hyper_parameters = checkpoint.get("hyper_parameters", None)
+        if hyper_parameters:
+            if ckpt_label_info := hyper_parameters.get("label_info"):
+                self._label_info = self._dispatch_label_info(ckpt_label_info)
+            if ckpt_tile_config := hyper_parameters.get("tile_config"):
+                if isinstance(ckpt_tile_config, dict):
+                    ckpt_tile_config = TileConfig(**ckpt_tile_config)
+                self.tile_config = ckpt_tile_config
 
     def load_state_dict_incrementally(self, ckpt: dict[str, Any], *args, **kwargs) -> None:
         """Load state dict incrementally."""
         ckpt_label_info: LabelInfo | None = (
-            ckpt.get("label_info") if not is_ckpt_from_otx_v1(ckpt) else self.get_ckpt_label_info_v1(ckpt)
+            ckpt.get("hyper_parameters", {}).get("label_info")
+            if not is_ckpt_from_otx_v1(ckpt)
+            else self.get_ckpt_label_info_v1(ckpt)
         )
 
         if ckpt_label_info is None:
             msg = "Checkpoint should have `label_info`."
             raise ValueError(msg, ckpt_label_info)
 
+        if isinstance(ckpt_label_info, dict):
+            ckpt_label_info = LabelInfo(**ckpt_label_info)
+
         if not hasattr(ckpt_label_info, "label_ids"):
             msg = "Loading checkpoint from OTX < 2.2.1, label_ids are assigned automatically"
             logger.info(msg)
@@ -447,10 +452,10 @@ def load_state_dict(self, ckpt: dict[str, Any], *args, **kwargs) -> None:
             warnings.warn(msg, stacklevel=2)
             state_dict = self.load_from_otx_v1_ckpt(ckpt)
         elif is_ckpt_for_finetuning(ckpt):
+            self.on_load_checkpoint(ckpt)
             state_dict = ckpt["state_dict"]
         else:
             state_dict = ckpt
-
         return super().load_state_dict(state_dict, *args, **kwargs)
 
     def load_from_otx_v1_ckpt(self, ckpt: dict[str, Any]) -> dict:
@@ -828,6 +833,11 @@ def get_dummy_input(self, batch_size: int = 1) -> OTXBatchDataEntity:
 
     @staticmethod
     def _dispatch_label_info(label_info: LabelInfoTypes) -> LabelInfo:
+        if isinstance(label_info, dict):
+            if "label_ids" not in label_info:
+                # NOTE: This is for backward compatibility
+                label_info["label_ids"] = label_info["label_names"]
+            return LabelInfo(**label_info)
         if isinstance(label_info, int):
             return LabelInfo.from_num_classes(num_classes=label_info)
         if isinstance(label_info, Sequence) and all(isinstance(name, str) for name in label_info):
@@ -837,6 +847,9 @@ def _dispatch_label_info(label_info: LabelInfoTypes) -> LabelInfo:
                 label_ids=[str(i) for i in range(len(label_info))],
             )
         if isinstance(label_info, LabelInfo):
+            if not hasattr(label_info, "label_ids"):
+                # NOTE: This is for backward compatibility
+                label_info.label_ids = label_info.label_names
             return label_info
 
         raise TypeError(label_info)

src/otx/core/model/classification.py

@@ -478,10 +478,18 @@ def _convert_pred_entity_to_compute_metric(
 
     @staticmethod
     def _dispatch_label_info(label_info: LabelInfoTypes) -> LabelInfo:
-        if not isinstance(label_info, HLabelInfo):
-            raise TypeError(label_info)
-
-        return label_info
+        if isinstance(label_info, dict):
+            if "label_ids" not in label_info:
+                # NOTE: This is for backward compatibility
+                label_info["label_ids"] = label_info["label_names"]
+            return HLabelInfo(**label_info)
+        if isinstance(label_info, HLabelInfo):
+            if not hasattr(label_info, "label_ids"):
+                # NOTE: This is for backward compatibility
+                label_info.label_ids = label_info.label_names
+            return label_info
+
+        raise TypeError(label_info)
 
     def get_dummy_input(self, batch_size: int = 1) -> HlabelClsBatchDataEntity:
         """Returns a dummy input for classification OV model."""

src/otx/core/model/instance_segmentation.py

@@ -283,6 +283,7 @@ def _export_parameters(self) -> TaskLevelExportParameters:
         # Instance segmentation needs to add empty label to satisfy MAPI wrapper requirements
         modified_label_info.label_names.insert(0, "otx_empty_lbl")
         modified_label_info.label_ids.insert(0, "None")
+        modified_label_info.label_groups[0].insert(0, "otx_empty_lbl")
 
         return super()._export_parameters.wrap(
             model_type="MaskRCNN",
@@ -762,7 +763,6 @@ def _log_metrics(self, meter: Metric, key: Literal["val", "test"], **compute_kwa
 
     def _create_label_info_from_ov_ir(self) -> LabelInfo:
         ov_model = self.model.get_model()
-
         if ov_model.has_rt_info(["model_info", "label_info"]):
             serialized = ov_model.get_rt_info(["model_info", "label_info"]).value
             ir_label_info = LabelInfo.from_json(serialized)

src/otx/core/model/segmentation.py

@@ -193,6 +193,11 @@ def _convert_pred_entity_to_compute_metric(
 
     @staticmethod
     def _dispatch_label_info(label_info: LabelInfoTypes) -> LabelInfo:
+        if isinstance(label_info, dict):
+            if "label_ids" not in label_info:
+                # NOTE: This is for backward compatibility
+                label_info["label_ids"] = label_info["label_names"]
+            return SegLabelInfo(**label_info)
         if isinstance(label_info, int):
             return SegLabelInfo.from_num_classes(num_classes=label_info)
         if isinstance(label_info, Sequence) and all(isinstance(name, str) for name in label_info):
@@ -202,6 +207,9 @@ def _dispatch_label_info(label_info: LabelInfoTypes) -> LabelInfo:
                 label_ids=[str(i) for i in range(len(label_info))],
             )
         if isinstance(label_info, SegLabelInfo):
+            if not hasattr(label_info, "label_ids"):
+                # NOTE: This is for backward compatibility
+                label_info.label_ids = label_info.label_names
             return label_info
 
         raise TypeError(label_info)

src/otx/core/utils/utils.py

@@ -28,7 +28,7 @@ def is_ckpt_from_otx_v1(ckpt: dict) -> bool:
     Returns:
         bool: True means the checkpoint comes from otx1
     """
-    return "model" in ckpt and ckpt["VERSION"] == 1
+    return "model" in ckpt and "VERSION" in ckpt and ckpt["VERSION"] == 1
 
 
 def is_ckpt_for_finetuning(ckpt: dict) -> bool: