Add deprecation tests on offsets models

CeliaBenquet · CeliaBenquet · commit dbae49a6baf1 · 2026-05-07T16:05:10.000+02:00
diff --git a/tests/_reference_implementations/reference_offset_models.py b/tests/_reference_implementations/reference_offset_models.py
@@ -0,0 +1,163 @@
+"""Reference implementations of previously hardcoded offset models before parametrization.
+
+These models are used to verify that the parametrized versions produce identical outputs.
+"""
+
+import torch
+from torch import nn
+
+import cebra.data
+import cebra.data.datatypes
+import cebra.models.layers as cebra_layers
+from cebra.models.model import _OffsetModel
+from cebra.models.model import ConvolutionalModelMixin
+
+
+class Offset10ModelReference(_OffsetModel, ConvolutionalModelMixin):
+    """Reference: CEBRA model with a 10 sample receptive field (offset10-model)."""
+
+    def __init__(self, num_neurons, num_units, num_output, normalize=True):
+        if num_units < 1:
+            raise ValueError(
+                f"Hidden dimension needs to be at least 1, but got {num_units}."
+            )
+        super().__init__(
+            nn.Conv1d(num_neurons, num_units, 2),
+            nn.GELU(),
+            *self._make_layers(num_units, num_layers=3),
+            nn.Conv1d(num_units, num_output, 3),
+            num_input=num_neurons,
+            num_output=num_output,
+            normalize=normalize,
+        )
+
+    def get_offset(self) -> cebra.data.datatypes.Offset:
+        return cebra.data.Offset(5, 5)
+
+
+class Offset5ModelReference(_OffsetModel, ConvolutionalModelMixin):
+    """Reference: CEBRA model with a 5 sample receptive field (offset5-model)."""
+
+    def __init__(self, num_neurons, num_units, num_output, normalize=True):
+        super().__init__(
+            nn.Conv1d(num_neurons, num_units, 2),
+            nn.GELU(),
+            cebra_layers._Skip(nn.Conv1d(num_units, num_units, 3), nn.GELU()),
+            nn.Conv1d(num_units, num_output, 2),
+            num_input=num_neurons,
+            num_output=num_output,
+            normalize=normalize,
+        )
+
+    def get_offset(self) -> cebra.data.datatypes.Offset:
+        return cebra.data.Offset(2, 3)
+
+
+class Offset15ModelReference(_OffsetModel, ConvolutionalModelMixin):
+    """Reference: CEBRA model with a 15 sample receptive field (offset15-model)."""
+
+    def __init__(self, num_neurons, num_units, num_output, normalize=True):
+        if num_units < 1:
+            raise ValueError(
+                f"Hidden dimension needs to be at least 1, but got {num_units}."
+            )
+        super().__init__(
+            nn.Conv1d(num_neurons, num_units, 2),
+            nn.GELU(),
+            *self._make_layers(num_units, num_layers=6),
+            nn.Conv1d(num_units, num_output, 2),
+            num_input=num_neurons,
+            num_output=num_output,
+            normalize=normalize,
+        )
+
+    def get_offset(self) -> cebra.data.datatypes.Offset:
+        return cebra.data.Offset(7, 8)
+
+
+class Offset20ModelReference(_OffsetModel, ConvolutionalModelMixin):
+    """Reference: CEBRA model with a 20 sample receptive field (offset20-model)."""
+
+    def __init__(self, num_neurons, num_units, num_output, normalize=True):
+        if num_units < 1:
+            raise ValueError(
+                f"Hidden dimension needs to be at least 1, but got {num_units}."
+            )
+        super().__init__(
+            nn.Conv1d(num_neurons, num_units, 2),
+            nn.GELU(),
+            *self._make_layers(num_units, num_layers=8),
+            nn.Conv1d(num_units, num_output, 3),
+            num_input=num_neurons,
+            num_output=num_output,
+            normalize=normalize,
+        )
+
+    def get_offset(self) -> cebra.data.datatypes.Offset:
+        return cebra.data.Offset(10, 10)
+
+
+class Offset36Reference(_OffsetModel, ConvolutionalModelMixin):
+    """Reference: CEBRA model with a 36 sample receptive field (offset36-model)."""
+
+    def __init__(self, num_neurons, num_units, num_output, normalize=True):
+        if num_units < 1:
+            raise ValueError(
+                f"Hidden dimension needs to be at least 1, but got {num_units}."
+            )
+        super().__init__(
+            nn.Conv1d(num_neurons, num_units, 2),
+            nn.GELU(),
+            *self._make_layers(num_units, num_layers=16),
+            nn.Conv1d(num_units, num_output, 3),
+            num_input=num_neurons,
+            num_output=num_output,
+            normalize=normalize,
+        )
+
+    def get_offset(self) -> cebra.data.datatypes.Offset:
+        return cebra.data.Offset(18, 18)
+
+
+class Offset40Reference(_OffsetModel, ConvolutionalModelMixin):
+    """Reference: CEBRA model with a 40 sample receptive field (offset40-model)."""
+
+    def __init__(self, num_neurons, num_units, num_output, normalize=True):
+        if num_units < 1:
+            raise ValueError(
+                f"Hidden dimension needs to be at least 1, but got {num_units}."
+            )
+        super().__init__(
+            nn.Conv1d(num_neurons, num_units, 2),
+            nn.GELU(),
+            *self._make_layers(num_units, 18),
+            nn.Conv1d(num_units, num_output, 3),
+            num_input=num_neurons,
+            num_output=num_output,
+            normalize=normalize,
+        )
+
+    def get_offset(self) -> cebra.data.datatypes.Offset:
+        return cebra.data.Offset(20, 20)
+
+
+class Offset50Reference(_OffsetModel, ConvolutionalModelMixin):
+    """Reference: CEBRA model with a 50 sample receptive field (offset50-model)."""
+
+    def __init__(self, num_neurons, num_units, num_output, normalize=True):
+        if num_units < 1:
+            raise ValueError(
+                f"Hidden dimension needs to be at least 1, but got {num_units}."
+            )
+        super().__init__(
+            nn.Conv1d(num_neurons, num_units, 2),
+            nn.GELU(),
+            *self._make_layers(num_units, 23),
+            nn.Conv1d(num_units, num_output, 3),
+            num_input=num_neurons,
+            num_output=num_output,
+            normalize=normalize,
+        )
+
+    def get_offset(self) -> cebra.data.datatypes.Offset:
+        return cebra.data.Offset(25, 25)
diff --git a/tests/test_models.py b/tests/test_models.py
@@ -165,3 +165,204 @@ def test_version_check_dropout_available():
         assert len(cebra.models.get_options("*dropout*")) == 0
     else:
         assert len(cebra.models.get_options("*dropout*")) > 0
+
+
+# Tests for parametrized offset models backward compatibility
+from _reference_implementations import Offset5ModelReference
+from _reference_implementations import Offset10ModelReference
+from _reference_implementations import Offset15ModelReference
+from _reference_implementations import Offset20ModelReference
+from _reference_implementations import Offset36Reference
+from _reference_implementations import Offset40Reference
+from _reference_implementations import Offset50Reference
+
+
+@pytest.mark.parametrize("offset_n,reference_class", [
+    (5, Offset5ModelReference),
+    (10, Offset10ModelReference),
+    (15, Offset15ModelReference),
+    (20, Offset20ModelReference),
+    (36, Offset36Reference),
+    (40, Offset40Reference),
+    (50, Offset50Reference),
+])
+def test_parametrized_offset_models_match_reference(offset_n, reference_class):
+    """Test that parametrized offset models produce identical output to reference hardcoded models."""
+
+    num_neurons = 5
+    num_units = 8
+    num_output = 3
+    normalize = True
+
+    # Create reference model
+    ref_model = reference_class(num_neurons,
+                                num_units,
+                                num_output,
+                                normalize=normalize)
+
+    # Create parametrized model using OffsetNModel
+    param_model = cebra.models.init(f"offset{offset_n}-model",
+                                    num_neurons=num_neurons,
+                                    num_units=num_units,
+                                    num_output=num_output)
+
+    # Test 1: Check offsets match
+    ref_offset = ref_model.get_offset()
+    param_offset = param_model.get_offset()
+    assert ref_offset.left == param_offset.left, \
+        f"Offset left mismatch for offset{offset_n}: {ref_offset.left} != {param_offset.left}"
+    assert ref_offset.right == param_offset.right, \
+        f"Offset right mismatch for offset{offset_n}: {ref_offset.right} != {param_offset.right}"
+
+    # Test 2: Check model architecture - same number of parameters
+    ref_params = sum(p.numel() for p in ref_model.parameters())
+    param_params = sum(p.numel() for p in param_model.parameters())
+    assert ref_params == param_params, \
+        f"Parameter count mismatch for offset{offset_n}: {ref_params} != {param_params}"
+
+    # Test 3: Check output shape consistency
+    batch_size = 2
+    input_length = 100
+    offset_len = len(ref_offset)
+
+    test_input = torch.randn(batch_size, num_neurons, offset_len)
+
+    with torch.no_grad():
+        ref_output = ref_model.net(test_input)
+        param_output = param_model.net(test_input)
+
+    assert ref_output.shape == param_output.shape, \
+        f"Output shape mismatch for offset{offset_n}: {ref_output.shape} != {param_output.shape}"
+
+    # Test 4: For convolutional models, test on full length input
+    if isinstance(param_model, cebra.models.ConvolutionalModelMixin):
+        test_input_full = torch.randn(batch_size, num_neurons, input_length)
+
+        with torch.no_grad():
+            ref_output_full = ref_model.net(test_input_full)
+            param_output_full = param_model.net(test_input_full)
+
+        expected_length = input_length - len(ref_offset) + 1
+        assert ref_output_full.shape == (batch_size, num_output, expected_length), \
+            f"Reference model output shape unexpected for offset{offset_n}"
+        assert param_output_full.shape == (batch_size, num_output, expected_length), \
+            f"Parametrized model output shape unexpected for offset{offset_n}"
+
+
+@pytest.mark.parametrize("offset_n", [5, 10, 15, 18, 20, 31, 36, 40, 50])
+def test_parametrized_offset_models_exist(offset_n):
+    """Test that all parametrized offset models can be instantiated."""
+    model = cebra.models.init(f"offset{offset_n}-model",
+                              num_neurons=5,
+                              num_units=4,
+                              num_output=3)
+    assert isinstance(model, cebra.models.Model)
+    assert isinstance(model, cebra.models.HasFeatureEncoder)
+    assert isinstance(model, cebra.models.ConvolutionalModelMixin)
+
+
+@pytest.mark.parametrize("offset_n,reference_class", [
+    (5, Offset5ModelReference),
+    (10, Offset10ModelReference),
+    (15, Offset15ModelReference),
+    (20, Offset20ModelReference),
+    (36, Offset36Reference),
+    (40, Offset40Reference),
+    (50, Offset50Reference),
+])
+def test_parametrized_offset_models_forward_pass_identical(
+        offset_n, reference_class):
+    """Test that parametrized and reference models produce identical forward pass outputs.
+    
+    This test verifies that when both models are initialized with the same seed and weights,
+    they produce identical outputs.
+    """
+
+    num_neurons = 5
+    num_units = 8
+    num_output = 3
+    normalize = True
+    batch_size = 2
+
+    # Set seed for reproducibility
+    torch.manual_seed(42)
+
+    # Create reference model and get its state dict
+    ref_model = reference_class(num_neurons,
+                                num_units,
+                                num_output,
+                                normalize=normalize)
+    ref_state_dict = {k: v.clone() for k, v in ref_model.state_dict().items()}
+
+    # Create parametrized model
+    param_model = cebra.models.init(f"offset{offset_n}-model",
+                                    num_neurons=num_neurons,
+                                    num_units=num_units,
+                                    num_output=num_output)
+
+    # Load the same weights into parametrized model
+    param_model.load_state_dict(ref_state_dict)
+
+    # Test with multiple input sizes
+    offset = ref_model.get_offset()
+    offset_len = len(offset)
+
+    for input_length in [offset_len, offset_len * 2, 100]:
+        test_input = torch.randn(batch_size, num_neurons, input_length)
+
+        with torch.no_grad():
+            ref_output = ref_model.net(test_input)
+            param_output = param_model.net(test_input)
+
+        # Check that outputs are identical
+        assert torch.allclose(ref_output, param_output, rtol=1e-5, atol=1e-7), \
+            f"Output mismatch for offset{offset_n} with input_length={input_length}"
+
+        # Check that outputs have same device and dtype
+        assert ref_output.device == param_output.device, \
+            f"Device mismatch for offset{offset_n}"
+        assert ref_output.dtype == param_output.dtype, \
+            f"Dtype mismatch for offset{offset_n}"
+
+
+@pytest.mark.parametrize("offset_n", [5, 10, 15, 18, 20, 31, 36, 40, 50])
+def test_parametrized_offset_models_layer_structure(offset_n):
+    """Test that parametrized models have the correct layer structure."""
+    num_neurons = 4
+    num_units = 8
+    num_output = 3
+
+    model = cebra.models.init(f"offset{offset_n}-model",
+                              num_neurons=num_neurons,
+                              num_units=num_units,
+                              num_output=num_output)
+
+    # Model should have Conv1d -> GELU -> Skip layers -> Conv1d structure
+    # Extract the actual network layers
+    layers = list(model.net.children())
+
+    # First layer should be Conv1d
+    assert isinstance(layers[0], nn.Conv1d), \
+        f"First layer of offset{offset_n} model should be Conv1d"
+    assert layers[0].in_channels == num_neurons
+    assert layers[0].out_channels == num_units
+    assert layers[0].kernel_size == (2,)
+
+    # Last meaningful layer (before Norm and Squeeze) should be Conv1d
+    # Find the second-to-last Conv1d layer
+    conv_layers = [l for l in layers if isinstance(l, nn.Conv1d)]
+    assert len(conv_layers) >= 2, \
+        f"offset{offset_n} model should have at least 2 Conv1d layers"
+
+    last_conv = conv_layers[-1]
+    assert last_conv.out_channels == num_output
+
+    # Check that offset is computed correctly
+    offset = model.get_offset()
+    expected_left = offset_n // 2
+    expected_right = offset_n // 2 + offset_n % 2
+
+    assert offset.left == expected_left, \
+        f"Offset left for offset{offset_n} should be {expected_left}, got {offset.left}"
+    assert offset.right == expected_right, \
+        f"Offset right for offset{offset_n} should be {expected_right}, got {offset.right}"
