From 5a8901757a8ad4aedc8ebd1d8371fcb80590ebe6 Mon Sep 17 00:00:00 2001 From: myscience Date: Thu, 22 Oct 2026 23:31:06 +0200 Subject: [PATCH 1/2] Preserve DCL dynamics parameters in solver checkpoints --- cebra/solver/single_session.py | 31 +++++++++++++- docs/source/api/pytorch/dynamics.rst | 12 +++++- tests/test_dcl_checkpoint.py | 62 ++++++++++++++++++++++++++++ 3 files changed, 102 insertions(+), 3 deletions(-) create mode 100644 tests/test_dcl_checkpoint.py 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) From 9db60a88e27e1495a3e9add52aa0bd7d170cdf36 Mon Sep 17 00:00:00 2001 From: myscience Date: Wed, 28 Oct 2026 18:39:04 +0100 Subject: [PATCH 2/2] Save reference models in the shared auxiliary solver --- cebra/solver/base.py | 28 +++++ cebra/solver/single_session.py | 31 +----- docs/source/api/pytorch/dynamics.rst | 14 ++- tests/test_auxiliary_checkpoint.py | 161 +++++++++++++++++++++++++++ tests/test_dcl_checkpoint.py | 62 ----------- 5 files changed, 199 insertions(+), 97 deletions(-) create mode 100644 tests/test_auxiliary_checkpoint.py delete mode 100644 tests/test_dcl_checkpoint.py diff --git a/cebra/solver/base.py b/cebra/solver/base.py index 2766cd79..ec64888f 100644 --- a/cebra/solver/base.py +++ b/cebra/solver/base.py @@ -892,6 +892,34 @@ def step(self, batch: cebra.data.Batch) -> dict: class AuxiliaryVariableSolver(Solver): + def state_dict(self) -> dict: + """Save the solver state, including the reference model parameters.""" + state = super().state_dict() + state["reference_model"] = self.reference_model.state_dict() + return state + + def load_state_dict(self, state_dict: dict, strict: bool = True): + """Restore the solver and reference model from a checkpoint. + + Args: + state_dict: Solver checkpoint, including ``reference_model``. + strict: Require reference model parameters. If ``False``, load + available components and warn when the reference model is absent. + """ + if "reference_model" not in state_dict: + message = ( + "Checkpoint is missing reference_model parameters. " + "Legacy auxiliary solver checkpoints did not save them. " + "Use strict=False only to load available components and leave " + "unsaved reference model parameters unchanged.") + if strict: + raise KeyError(message) + warnings.warn(message, UserWarning, stacklevel=2) + else: + self.reference_model.load_state_dict(state_dict["reference_model"], + strict=strict) + super().load_state_dict(state_dict, strict=strict) + @torch.no_grad() def transform(self, inputs: torch.Tensor, diff --git a/cebra/solver/single_session.py b/cebra/solver/single_session.py index 784b9084..872f9f8e 100644 --- a/cebra/solver/single_session.py +++ b/cebra/solver/single_session.py @@ -22,7 +22,6 @@ """Single session solvers embed a single pair of time series.""" import copy -import warnings from typing import Optional import literate_dataclasses as dataclasses @@ -225,35 +224,9 @@ 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(_DynamicsStateMixin, - SingleSessionAuxVariableSolver): +class SingleSessionDCLSolver(SingleSessionAuxVariableSolver): dynamics_model: torch.nn.Module = None def __post_init__(self): @@ -395,7 +368,7 @@ def _inference(self, batch: cebra.data.Batch) -> cebra.data.Batch: @register("single-session-full-dcl") @dataclasses.dataclass -class BatchSingleSessionDCLSolver(_DynamicsStateMixin, BatchSingleSessionSolver, +class BatchSingleSessionDCLSolver(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 d34699d6..3e65c463 100644 --- a/docs/source/api/pytorch/dynamics.rst +++ b/docs/source/api/pytorch/dynamics.rst @@ -28,9 +28,11 @@ Models 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. +Auxiliary solver checkpoints include ``reference_model``. For both minibatch +and full-batch DCL solvers, this contains the encoder and dynamics parameters. +Historical auxiliary solver checkpoints omitted the reference model and cannot +restore its unsaved parameters, including trained DCL dynamics. Strict loading +raises an explanatory error. Low-level +``solver.load_state_dict(checkpoint, strict=False)`` loads the available +components and warns that unsaved reference model parameters retain their +current values. Solvers without an auxiliary reference model are unchanged. diff --git a/tests/test_auxiliary_checkpoint.py b/tests/test_auxiliary_checkpoint.py new file mode 100644 index 00000000..836c39cc --- /dev/null +++ b/tests/test_auxiliary_checkpoint.py @@ -0,0 +1,161 @@ +# +# 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. +# +"""Auxiliary solver checkpoints preserve reference models and DCL dynamics.""" + +import copy +import itertools + +import numpy as np +import pytest +import torch + +import cebra +import cebra.models +import cebra.solver + + +def _make_auxiliary_solver(multisession): + + def make_encoder(): + return torch.nn.Sequential(torch.nn.Linear(5, 3), + torch.nn.BatchNorm1d(3)) + + if multisession: + model = torch.nn.ModuleList([make_encoder(), make_encoder()]) + reference_model = torch.nn.ModuleList([make_encoder(), make_encoder()]) + solver_class = cebra.solver.MultiSessionAuxVariableSolver + else: + model = make_encoder() + reference_model = make_encoder() + solver_class = cebra.solver.SingleSessionAuxVariableSolver + criterion = cebra.models.InfoNCE() + optimizer = torch.optim.Adam( + itertools.chain(model.parameters(), reference_model.parameters(), + criterion.parameters())) + kwargs = {} if multisession else {"reference_model": reference_model} + solver = solver_class(model=model, + criterion=criterion, + optimizer=optimizer, + tqdm_on=False, + **kwargs) + if multisession: + solver.reference_model = reference_model + return solver + + +@pytest.mark.parametrize("multisession", [False, True]) +def test_auxiliary_reference_round_trip(multisession): + torch.manual_seed(4) + solver = _make_auxiliary_solver(multisession) + restored = _make_auxiliary_solver(multisession) + inputs = torch.randn(8, 5) + models = list(solver.model) if multisession else [solver.model] + references = (list(solver.reference_model) + if multisession else [solver.reference_model]) + for reference in references: + reference(inputs) # Update the reference encoder's running statistics. + solver.model.eval() + solver.reference_model.eval() + restored.model.eval() + restored.reference_model.eval() + restored_references = (list(restored.reference_model) + if multisession else [restored.reference_model]) + for reference, other in zip(references, restored_references): + assert not torch.allclose(reference(inputs), other(inputs)) + + restored.load_state_dict(copy.deepcopy(solver.state_dict())) + restored_models = list(restored.model) if multisession else [restored.model] + for model, other in zip(models, restored_models): + torch.testing.assert_close(model(inputs), other(inputs)) + for reference, other in zip(references, restored_references): + torch.testing.assert_close(reference(inputs), other(inputs)) + for name, value in reference.named_buffers(): + torch.testing.assert_close(value, dict(other.named_buffers())[name]) + + +@pytest.mark.parametrize("multisession", [False, True]) +def test_auxiliary_legacy_checkpoint(multisession): + solver = _make_auxiliary_solver(multisession) + restored = _make_auxiliary_solver(multisession) + legacy = copy.deepcopy(solver.state_dict()) + del legacy["reference_model"] + before = copy.deepcopy(restored.reference_model.state_dict()) + with pytest.raises(KeyError, match="missing reference_model"): + restored.load_state_dict(legacy) + with pytest.warns(UserWarning, match="Legacy auxiliary"): + restored.load_state_dict(legacy, strict=False) + for name, value in before.items(): + torch.testing.assert_close(value, + restored.reference_model.state_dict()[name]) + for name, value in solver.model.state_dict().items(): + torch.testing.assert_close(value, restored.model.state_dict()[name]) + + +@pytest.mark.parametrize("strict", [False, True]) +def test_auxiliary_reference_parameter_strictness(strict): + solver = _make_auxiliary_solver(multisession=False) + checkpoint = copy.deepcopy(solver.state_dict()) + del checkpoint["reference_model"]["0.weight"] + if strict: + with pytest.raises(RuntimeError, match="Missing key.*0.weight"): + solver.load_state_dict(checkpoint) + else: + solver.load_state_dict(checkpoint, strict=False) + + +@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") + assert loaded.solver_.reference_model[0] is loaded.solver_.model + assert loaded.solver_.reference_model[1] is loaded.dynamics_model_ + 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["reference_model"] + before = loaded.dynamics_model_(state).detach().clone() + with pytest.raises(KeyError, match="missing reference_model"): + loaded.solver_.load_state_dict(legacy) + with pytest.warns(UserWarning, match="Legacy auxiliary"): + loaded.solver_.load_state_dict(legacy, strict=False) + torch.testing.assert_close(loaded.dynamics_model_(state), before) + np.testing.assert_allclose(model.transform(data), loaded.transform(data)) diff --git a/tests/test_dcl_checkpoint.py b/tests/test_dcl_checkpoint.py deleted file mode 100644 index 34481b00..00000000 --- a/tests/test_dcl_checkpoint.py +++ /dev/null @@ -1,62 +0,0 @@ -# -# 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)