Enable torch.complile to work with classification

src/otx/core/model/base.py

@@ -51,7 +51,7 @@
 from otx.core.types.precision import OTXPrecisionType
 from otx.core.utils.build import get_default_num_async_infer_requests
 from otx.core.utils.miscellaneous import ensure_callable
-from otx.core.utils.utils import is_ckpt_for_finetuning, is_ckpt_from_otx_v1
+from otx.core.utils.utils import is_ckpt_for_finetuning, is_ckpt_from_otx_v1, remove_state_dict_prefix
 
 if TYPE_CHECKING:
     from pathlib import Path
@@ -353,6 +353,10 @@ def _log_metrics(self, meter: Metric, key: Literal["val", "test"], **compute_kwa
 
     def on_save_checkpoint(self, checkpoint: dict[str, Any]) -> None:
         """Callback on saving checkpoint."""
+        if self.torch_compile:
+            # If torch_compile is True, a prefix key named _orig_mod. is added to the state_dict. Remove this.
+            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

src/otx/core/utils/utils.py

@@ -92,3 +92,12 @@ def import_object_from_module(obj_path: str) -> Any:  # noqa: ANN401
     module_name, obj_name = obj_path.rsplit(".", 1)
     module = importlib.import_module(module_name)
     return getattr(module, obj_name)
+
+
+def remove_state_dict_prefix(state_dict: dict[str, Any], prefix: str) -> dict[str, Any]:
+    """Remove prefix from state_dict keys."""
+    new_state_dict = {}
+    for key, value in state_dict.items():
+        new_key = key.replace(prefix, "")
+        new_state_dict[new_key] = value
+    return new_state_dict

tests/unit/core/utils/test_utils.py

@@ -12,6 +12,7 @@
     import_object_from_module,
     is_ckpt_for_finetuning,
     is_ckpt_from_otx_v1,
+    remove_state_dict_prefix,
 )
 
 
@@ -120,3 +121,20 @@ def test_import_object_from_module():
     obj_path = "otx.core.utils.utils.get_mean_std_from_data_processing"
     obj = import_object_from_module(obj_path)
     assert obj == get_mean_std_from_data_processing
+
+
+def test_remove_state_dict_prefix():
+    state_dict = {
+        "model._orig_mod.backbone.0.weight": 1,
+        "model._orig_mod.backbone.0.bias": 2,
+        "model._orig_mod.backbone.1.weight": 3,
+        "model._orig_mod.backbone.1.bias": 4,
+    }
+    new_state_dict = remove_state_dict_prefix(state_dict=state_dict, prefix="_orig_mod.")
+    expected = {
+        "model.backbone.0.weight": 1,
+        "model.backbone.0.bias": 2,
+        "model.backbone.1.weight": 3,
+        "model.backbone.1.bias": 4,
+    }
+    assert new_state_dict == expected