Fix legacy loading logic

Author: stes

cebra/integrations/sklearn/cebra.py

@@ -23,6 +23,8 @@
 
 import importlib.metadata
 import itertools
+import pickle
+import warnings
 from typing import (Callable, Dict, Iterable, List, Literal, Optional, Tuple,
                     Union)
 
@@ -1336,6 +1338,26 @@ def _get_state(self):
         }
         return state
 
+    def _get_state_dict(self):
+        backend = "sklearn"
+        return {
+            'args': self.get_params(),
+            'state': self._get_state(),
+            'state_dict': self.solver_.state_dict(),
+            'metadata': {
+                'backend':
+                    backend,
+                'cebra_version':
+                    cebra.__version__,
+                'torch_version':
+                    torch.__version__,
+                'numpy_version':
+                    np.__version__,
+                'sklearn_version':
+                    importlib.metadata.distribution("scikit-learn").version
+            }
+        }
+
     def save(self,
              filename: str,
              backend: Literal["torch", "sklearn"] = "sklearn"):
@@ -1384,28 +1406,16 @@ def save(self,
         """
         if sklearn_utils.check_fitted(self):
             if backend == "torch":
+                warnings.warn(
+                    "Saving with backend='torch' is deprecated and will be removed in a future version. "
+                    "Please use backend='sklearn' instead.",
+                    DeprecationWarning,
+                    stacklevel=2,
+                )
                 checkpoint = torch.save(self, filename)
 
             elif backend == "sklearn":
-                checkpoint = torch.save(
-                    {
-                        'args': self.get_params(),
-                        'state': self._get_state(),
-                        'state_dict': self.solver_.state_dict(),
-                        'metadata': {
-                            'backend':
-                                backend,
-                            'cebra_version':
-                                cebra.__version__,
-                            'torch_version':
-                                torch.__version__,
-                            'numpy_version':
-                                np.__version__,
-                            'sklearn_version':
-                                importlib.metadata.distribution("scikit-learn"
-                                                               ).version
-                        }
-                    }, filename)
+                checkpoint = torch.save(self._get_state_dict(), filename)
             else:
                 raise NotImplementedError(f"Unsupported backend: {backend}")
         else:
@@ -1457,15 +1467,52 @@ def load(cls,
             >>> tmp_file.unlink()
         """
         supported_backends = ["auto", "sklearn", "torch"]
+
         if backend not in supported_backends:
             raise NotImplementedError(
                 f"Unsupported backend: '{backend}'. Supported backends are: {', '.join(supported_backends)}"
             )
 
-        checkpoint = _safe_torch_load(filename, weights_only, **kwargs)
+        if backend not in ["auto", "sklearn"]:
+            warnings.warn(
+                "From CEBRA version 0.6.1 onwards, the 'backend' parameter in cebra.CEBRA.load is deprecated and will be ignored; "
+                "the sklearn backend is now always used. Models saved with the torch backend can still be loaded.",
+                category=DeprecationWarning,
+                stacklevel=2,
+            )
 
-        if backend == "auto":
-            backend = "sklearn" if isinstance(checkpoint, dict) else "torch"
+        backend = "sklearn"
+
+        # NOTE(stes): For maximum backwards compatibility, we allow to load legacy checkpoints. From 0.7.0 onwards,
+        # the user will have to explicitly pass weights_only=False to load these checkpoints, following the changes
+        # introduced in torch 2.6.0.
+        try:
+            checkpoint = _safe_torch_load(filename, weights_only=True, **kwargs)
+        except pickle.UnpicklingError as e:
+            if weights_only is not False:
+                if packaging.version.parse(
+                        cebra.__version__) < packaging.version.parse("0.7"):
+                    warnings.warn(
+                        "Failed to unpickle checkpoint with weights_only=True. "
+                        "Falling back to loading with weights_only=False. "
+                        "This is unsafe and should only be done if you trust the source of the model file. "
+                        "In the future, loading these checkpoints will only work if weights_only=False is explicitly passed.",
+                        category=UserWarning,
+                        stacklevel=2,
+                    )
+                else:
+                    raise ValueError(
+                        "Failed to unpickle checkpoint with weights_only=True. "
+                        "This may be due to an incompatible model file format. "
+                        "To attempt loading this checkpoint, please pass weights_only=False to CEBRA.load. "
+                        "Example: CEBRA.load(filename, weights_only=False)."
+                    ) from e
+
+            checkpoint = _safe_torch_load(filename,
+                                          weights_only=False,
+                                          **kwargs)
+            checkpoint = _check_type_checkpoint(checkpoint)
+            checkpoint = checkpoint._get_state_dict()
 
         if isinstance(checkpoint, dict) and backend == "torch":
             raise RuntimeError(
@@ -1476,10 +1523,10 @@ def load(cls,
                 "Cannot use 'sklearn' backend a non dictionary-based checkpoint. "
                 "Please try a different backend.")
 
-        if backend == "sklearn":
-            cebra_ = _load_cebra_with_sklearn_backend(checkpoint)
-        else:
-            cebra_ = _check_type_checkpoint(checkpoint)
+        if backend != "sklearn":
+            raise ValueError(f"Unsupported backend: {backend}")
+
+        cebra_ = _load_cebra_with_sklearn_backend(checkpoint)
 
         n_features = cebra_.n_features_
         cebra_.solver_.n_features = ([

cebra/registry.py

@@ -46,6 +46,7 @@
 from __future__ import annotations
 
 import fnmatch
+import functools
 import itertools
 import sys
 import textwrap
@@ -214,14 +215,29 @@ def _zip_dict(d):
                 yield dict(zip(keys, combination))
 
         def _create_class(cls, **default_kwargs):
+            class_name = pattern.format(**default_kwargs)
 
-            @register(pattern.format(**default_kwargs), base=pattern)
+            @register(class_name, base=pattern)
             class _ParametrizedClass(cls):
 
                 def __init__(self, *args, **kwargs):
                     default_kwargs.update(kwargs)
                     super().__init__(*args, **default_kwargs)
 
+            # Make the class pickleable by copying metadata from the base class
+            # and registering it in the module namespace
+            functools.update_wrapper(_ParametrizedClass, cls, updated=[])
+
+            # Set a unique qualname so pickle finds this class, not the base class
+            unique_name = f"{cls.__qualname__}_{class_name.replace('-', '_')}"
+            _ParametrizedClass.__qualname__ = unique_name
+            _ParametrizedClass.__name__ = unique_name
+
+            # Register in module namespace so pickle can find it via getattr
+            parent_module = sys.modules.get(cls.__module__)
+            if parent_module is not None:
+                setattr(parent_module, unique_name, _ParametrizedClass)
+
         def _parametrize(cls):
             for _default_kwargs in kwargs:
                 _create_class(cls, **_default_kwargs)

tests/test_sklearn.py

@@ -1053,7 +1053,7 @@ def get_offset(self) -> cebra.data.datatypes.Offset:
 
 @pytest.mark.parametrize("action", _iterate_actions())
 @pytest.mark.parametrize("backend_save", ["torch", "sklearn"])
-@pytest.mark.parametrize("backend_load", ["auto", "torch", "sklearn"])
+@pytest.mark.parametrize("backend_load", ["sklearn", "auto", "torch"])
 @pytest.mark.parametrize("model_architecture",
                          ["offset1-model", "parametrized-model-5"])
 @pytest.mark.parametrize("device", ["cpu"] +
@@ -1072,20 +1072,14 @@ def test_save_and_load(action, backend_save, backend_load, model_architecture,
             with pytest.raises(ValueError):
                 original_model.save(tempname, backend=backend_save)
         else:
-            if "parametrized" in original_model.model_architecture and backend_save == "torch":
-                with pytest.raises(AttributeError):
-                    original_model.save(tempname, backend=backend_save)
-            else:
-                original_model.save(tempname, backend=backend_save)
+            original_model.save(tempname, backend=backend_save)
+
+            weights_only = None
 
-                if (backend_load != "auto") and (backend_save != backend_load):
-                    with pytest.raises(RuntimeError):
-                        cebra_sklearn_cebra.CEBRA.load(tempname, backend_load)
-                else:
-                    loaded_model = cebra_sklearn_cebra.CEBRA.load(
-                        tempname, backend_load)
-                    _assert_equal(original_model, loaded_model)
-                    action(loaded_model)
+            loaded_model = cebra_sklearn_cebra.CEBRA.load(
+                tempname, backend_load, weights_only=weights_only)
+            _assert_equal(original_model, loaded_model)
+            action(loaded_model)
 
 
 def get_ordered_cuda_devices():
@@ -1489,7 +1483,7 @@ def test_new_transform(model_architecture, device):
                                                               X,
                                                               session_id=0)
     assert np.allclose(embedding1, embedding2, rtol=1e-5,
-                       atol=1e-8), "Arrays are not close enough"
+                       atol=1e-8), " are not close enough"
 
     embedding1 = cebra_model.transform(torch.Tensor(X), session_id=0)
     embedding2 = _utils_deprecated.cebra_transform_deprecated(cebra_model,
@@ -1603,3 +1597,22 @@ def test_read_write():
         cebra_model.save(tempname)
         loaded_model = cebra.CEBRA.load(tempname)
         _assert_equal(cebra_model, loaded_model)
+
+
+def test_repro_pickle_error():
+    """The torch backend for save/loading fails with python 3.14.
+
+    See https://github.com/AdaptiveMotorControlLab/CEBRA/pull/292.
+
+    This test is a minimal repro of the error.
+    """
+
+    model = cebra_sklearn_cebra.CEBRA(model_architecture='parametrized-model-5',
+                                      max_iterations=5,
+                                      batch_size=100,
+                                      device='cpu')
+
+    model.fit(np.random.randn(1000, 10))
+
+    with _windows_compatible_tempfile(mode="w+b", delete=True) as tempname:
+        model.save(tempname, backend="torch")