diff --git a/cebra/solver/single_session.py b/cebra/solver/single_session.py index 872f9f8e..784b9084 100644 --- a/cebra/solver/single_session.py +++ b/cebra/solver/single_session.py @@ -22,6 +22,7 @@ """Single session solvers embed a single pair of time series.""" import copy +import warnings from typing import Optional import literate_dataclasses as dataclasses @@ -224,9 +225,35 @@ def _inference(self, batch: cebra.data.Batch) -> cebra.data.Batch: return cebra.data.Batch(ref, pos, neg) +class _DynamicsStateMixin: + """Include dynamics parameters in solver checkpoints.""" + + def state_dict(self) -> dict: + """Save the encoder, optimizer, criterion, and dynamics model.""" + state = super().state_dict() + state["dynamics_model"] = self.dynamics_model.state_dict() + return state + + def load_state_dict(self, state_dict: dict, strict: bool = True): + """Restore dynamics; reject incomplete legacy checkpoints in strict mode.""" + if "dynamics_model" not in state_dict: + message = ( + "Checkpoint is missing dynamics_model parameters. " + "Legacy DCL checkpoints did not save them. " + "Use strict=False only to retain the initialized dynamics.") + if strict: + raise KeyError(message) + warnings.warn(message, UserWarning, stacklevel=2) + else: + self.dynamics_model.load_state_dict(state_dict["dynamics_model"], + strict=strict) + super().load_state_dict(state_dict, strict=strict) + + @register("single-session-dcl") @dataclasses.dataclass -class SingleSessionDCLSolver(SingleSessionAuxVariableSolver): +class SingleSessionDCLSolver(_DynamicsStateMixin, + SingleSessionAuxVariableSolver): dynamics_model: torch.nn.Module = None def __post_init__(self): @@ -368,7 +395,7 @@ def _inference(self, batch: cebra.data.Batch) -> cebra.data.Batch: @register("single-session-full-dcl") @dataclasses.dataclass -class BatchSingleSessionDCLSolver(BatchSingleSessionSolver, +class BatchSingleSessionDCLSolver(_DynamicsStateMixin, BatchSingleSessionSolver, SingleSessionAuxVariableSolver): """Optimize a DCL model with batch gradient descent. diff --git a/docs/source/api/pytorch/dynamics.rst b/docs/source/api/pytorch/dynamics.rst index d3cd6fd2..d34699d6 100644 --- a/docs/source/api/pytorch/dynamics.rst +++ b/docs/source/api/pytorch/dynamics.rst @@ -23,4 +23,14 @@ Models .. automodule:: cebra.dynamics.linear :members: - :show-inheritance: \ No newline at end of file + :show-inheritance: + +Checkpoint behavior +~~~~~~~~~~~~~~~~~~~ + +Both minibatch and full-batch DCL solver checkpoints include the dynamics +parameters. Historical DCL checkpoints that omitted those parameters cannot +restore the trained dynamics. Strict loading raises an explanatory error. +Low-level ``solver.load_state_dict(checkpoint, strict=False)`` loads the +available components and warns that the dynamics retain their initialization. +Ordinary CEBRA checkpoints have no dynamics requirement. diff --git a/tests/test_dcl_checkpoint.py b/tests/test_dcl_checkpoint.py new file mode 100644 index 00000000..34481b00 --- /dev/null +++ b/tests/test_dcl_checkpoint.py @@ -0,0 +1,62 @@ +# +# CEBRA: Consistent EmBeddings of high-dimensional Recordings using Auxiliary variables +# © Mackenzie W. Mathis & Steffen Schneider (v0.4.0+) +# Source code: +# https://github.com/AdaptiveMotorControlLab/CEBRA +# +# Please see LICENSE.md for the full license document: +# https://github.com/AdaptiveMotorControlLab/CEBRA/blob/main/LICENSE.md +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""DCL checkpoints preserve dynamics as well as encoder parameters.""" + +import numpy as np +import pytest +import torch + +import cebra + + +@pytest.mark.parametrize("architecture", ["linear", "orthogonal-linear"]) +@pytest.mark.parametrize("batch_size", [8, None]) +def test_existing_dcl_fit_save_load_and_resume(tmp_path, architecture, + batch_size): + from cebra.integrations.sklearn.cebra import DCL + torch.manual_seed(4) + data = np.random.default_rng(4).normal(size=(100, 5)).astype("float32") + model = DCL(device="cpu", + dynamics_model_architecture=architecture, + output_dimension=3, + batch_size=batch_size, + batch_size_negatives=11, + max_iterations=2, + distance="euclidean") + model.fit(data) + path = tmp_path / "flow.pt" + model.save(path, backend="sklearn") + loaded = cebra.CEBRA.load(path, backend="sklearn") + np.testing.assert_allclose(model.transform(data), loaded.transform(data)) + state = torch.randn(7, 3) + torch.testing.assert_close(model.dynamics_model_(state), + loaded.dynamics_model_(state)) + before = loaded.dynamics_model_(state).detach().clone() + loaded.partial_fit(data) + assert torch.isfinite(loaded.dynamics_model_(state)).all() + assert not torch.equal(before, loaded.dynamics_model_(state)) + legacy = model.solver_.state_dict() + del legacy["dynamics_model"] + with pytest.raises(KeyError, match="missing dynamics_model"): + loaded.solver_.load_state_dict(legacy) + with pytest.warns(UserWarning, match="Legacy DCL"): + loaded.solver_.load_state_dict(legacy, strict=False)