Fix save/load

Author: CeliaBenquet

cebra/integrations/sklearn/cebra.py

@@ -1417,6 +1417,11 @@ def load(cls,
         else:
             cebra_ = _check_type_checkpoint(checkpoint)
 
+        n_features = cebra_.n_features_
+        cebra_.solver_.n_features = ([
+            session_n_features for session_n_features in n_features
+        ] if isinstance(n_features, list) else n_features)
+
         return cebra_
 
     def to(self, device: Union[str, torch.device]):

cebra/solver/base.py

@@ -633,13 +633,12 @@ def load(self, logdir, filename="checkpoint.pth"):
         checkpoint = torch.load(savepath, map_location=self.device)
         self.load_state_dict(checkpoint, strict=True)
 
-        if hasattr(self.model, "n_features"):
-            n_features = self.model.n_features
-            self.n_features = ([
-                session_n_features for session_n_features in n_features
-            ] if isinstance(n_features, list) else n_features)
+        n_features = self.n_features
+        self.n_features = ([
+            session_n_features for session_n_features in n_features
+        ] if isinstance(n_features, list) else n_features)
 
-    def save(self, logdir, filename="checkpoint_last.pth"):
+    def save(self, logdir, filename="checkpoint.pth"):
         """Save the model and optimizer params.
 
         Args:

tests/test_solver.py

@@ -19,7 +19,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+import copy
 import itertools
+import tempfile
 
 import numpy as np
 import pytest
@@ -91,6 +93,48 @@ def _make_model(dataset, model_architecture="offset10-model"):
 #     )
 
 
+def _assert_same_state_dict(first, second):
+    assert first.keys() == second.keys()
+    for key in first:
+        if isinstance(first[key], torch.Tensor):
+            assert torch.allclose(first[key], second[key]), key
+        elif isinstance(first[key], dict):
+            _assert_same_state_dict(first[key], second[key]), key
+        else:
+            assert first[key] == second[key]
+
+
+def check_if_fit(model):
+    """Check if a model was already fit.
+
+    Args:
+        model: The model to check.
+
+    Returns:
+        True if the model was already fit.
+    """
+    return hasattr(model, "n_features_")
+
+
+def _assert_equal(original_solver, loaded_solver):
+    for k in original_solver.model.state_dict():
+        assert original_solver.model.state_dict()[k].all(
+        ) == loaded_solver.model.state_dict()[k].all()
+    assert check_if_fit(loaded_solver) == check_if_fit(original_solver)
+
+    if check_if_fit(loaded_solver):
+        _assert_same_state_dict(original_solver.state_dict_,
+                                loaded_solver.state_dict_)
+        X = np.random.normal(0, 1, (100, 1))
+
+        if loaded_solver.num_sessions is not None:
+            assert np.allclose(loaded_solver.transform(X, session_id=0),
+                               original_solver.transform(X, session_id=0))
+        else:
+            assert np.allclose(loaded_solver.transform(X),
+                               original_solver.transform(X))
+
+
 @pytest.mark.parametrize(
     "data_name, loader_initfunc, model_architecture, solver_initfunc",
     single_session_tests)
@@ -144,6 +188,12 @@ def test_single_session(data_name, loader_initfunc, model_architecture,
     for param in solver.parameters():
         assert isinstance(param, torch.Tensor)
 
+    fitted_solver = copy.deepcopy(solver)
+    with tempfile.TemporaryDirectory() as temp_dir:
+        solver.save(temp_dir)
+        solver.load(temp_dir)
+    _assert_equal(fitted_solver, solver)
+
 
 @pytest.mark.parametrize(
     "data_name, loader_initfunc, model_architecture, solver_initfunc",
@@ -225,6 +275,12 @@ def test_single_session_hybrid(data_name, loader_initfunc, model_architecture,
     for param in solver.parameters():
         assert isinstance(param, torch.Tensor)
 
+    fitted_solver = copy.deepcopy(solver)
+    with tempfile.TemporaryDirectory() as temp_dir:
+        solver.save(temp_dir)
+        solver.load(temp_dir)
+    _assert_equal(fitted_solver, solver)
+
 
 @pytest.mark.parametrize(
     "data_name, loader_initfunc, model_architecture, solver_initfunc",
@@ -302,6 +358,12 @@ def test_multi_session(data_name, loader_initfunc, model_architecture,
         for param in solver.parameters():
             assert isinstance(param, torch.Tensor)
 
+    fitted_solver = copy.deepcopy(solver)
+    with tempfile.TemporaryDirectory() as temp_dir:
+        solver.save(temp_dir)
+        solver.load(temp_dir)
+    _assert_equal(fitted_solver, solver)
+
 
 @pytest.mark.parametrize(
     "inputs, add_padding, offset, start_batch_idx, end_batch_idx, expected_output",