Add legacy multiobjective model for backward compat

Author: stes

cebra/models/multiobjective.py

@@ -20,7 +20,7 @@
 # limitations under the License.
 #
 import itertools
-from typing import List
+from typing import List, Tuple
 
 import torch
 from torch import nn
@@ -106,6 +106,140 @@ def forward(self, inp):
         return inp / torch.norm(inp, dim=1, keepdim=True)
 
 
+class LegacyMultiobjectiveModel(nn.Module):
+    """Wrapper around contrastive learning models to all training with multiple objectives
+
+    Multi-objective training splits the last layer's feature representation into multiple
+    chunks, which are then used for individual training objectives.
+
+    Args:
+        module: The module to wrap
+        dimensions: A tuple of dimension values to extract from the model's feature embedding.
+        renormalize: If True, the individual feature slices will be re-normalized before
+            getting returned---this option only makes sense in conjunction with a loss based
+            on the cosine distance or dot product.
+        output_mode: A mode as defined in ``MultiobjectiveModel.Mode``. Overlapping means that
+            when ``dimensions`` are set to `(x0, x1, ...)``, features will be extracted from
+            ``0:x0, 0:x1, ...``. When mode is set to separate, features are extracted from
+            ``x0:x1, x1:x2, ...``.
+        append_last_dimension: Defaults to True, and will allow to omit the last dimension in
+            the ``dimensions`` argument (which should be equal to the output dimension) of the
+            given model.
+
+    TODO:
+        - Update nn.Module type annotation for ``module`` to cebra.models.Model
+    """
+
+    class Mode:
+        """Mode for slicing and potentially normalizing the output embedding.
+
+        The options are:
+
+        - ``OVERLAPPING``: When ``dimensions`` are set to `(x0, x1, ...)``, features will be
+          extracted from ``0:x0, 0:x1, ...``.
+        - ``SEPARATE``: Features are extracted from ``x0:x1, x1:x2, ...``
+
+        """
+
+        OVERLAPPING = "overlapping"
+        SEPARATE = "separate"
+        _ALL = {OVERLAPPING, SEPARATE}
+
+        def is_valid(self, mode):
+            """Check if a given string representation is valid.
+
+            Args:
+                mode: String representation of the mode.
+
+            Returns:
+                ``True`` for a valid representation, ``False`` otherwise.
+            """
+            return mode in _ALL  # noqa: F821
+
+    def __init__(
+        self,
+        module: nn.Module,
+        dimensions: Tuple[int],
+        renormalize: bool = False,
+        output_mode: str = "overlapping",
+        append_last_dimension: bool = False,
+    ):
+        super().__init__()
+
+        if not isinstance(module, cebra.models.Model):
+            raise ValueError("Can only wrap models that are subclassing the "
+                             "cebra.models.Model abstract base class. "
+                             f"Got a model of type {type(module)}.")
+
+        self.module = module
+        self.renormalize = renormalize
+        self.output_mode = output_mode
+
+        self._norm = _Norm()
+        self._compute_slices(dimensions, append_last_dimension)
+
+    @property
+    def get_offset(self):
+        """See :py:meth:`cebra.models.model.Model.get_offset`."""
+        return self.module.get_offset
+
+    @property
+    def num_output(self):
+        """See :py:attr:`cebra.models.model.Model.num_output`."""
+        return self.module.num_output
+
+    def _compute_slices(self, dimensions, append_last_dimension):
+
+        def _valid_dimensions(dimensions):
+            return max(dimensions) == self.num_output
+
+        if append_last_dimension:
+            if _valid_dimensions(dimensions):
+                raise ValueError(
+                    f"append_last_dimension should only be used if extra values are "
+                    f"available. Last requested dimensionality is already {dimensions[-1]}."
+                )
+            dimensions += (self.num_output,)
+        if not _valid_dimensions(dimensions):
+            raise ValueError(
+                f"Max of given dimensions needs to match the number of outputs "
+                f"in the encoder network. Got {dimensions} and expected a "
+                f"maximum value of {self.num_output}.")
+
+        if self.output_mode == self.Mode.OVERLAPPING:
+            self.feature_ranges = tuple(
+                slice(0, dimension) for dimension in dimensions)
+        elif self.output_mode == self.Mode.SEPARATE:
+            from_dimension = (0,) + dimensions
+            self.feature_ranges = tuple(
+                slice(i, j) for i, j in zip(from_dimension, dimensions))
+        else:
+            raise ValueError(
+                f"Unknown mode: '{self.output_mode}', use one of {self.Mode._ALL}."
+            )
+
+    def forward(self, inputs):
+        """Compute multiple embeddings for a single signal input.
+
+        Args:
+            inputs: The input tensor
+
+        Returns:
+            A tuple of tensors which are sliced according to `self.feature_ranges`
+            if `renormalize` is set to true, each of the tensors will be normalized
+            across the first (feature) dimension.
+
+        TODO:
+            - Cover this function with unit tests
+        """
+        output = self.module(inputs)
+        outputs = (
+            output[:, slice_features] for slice_features in self.feature_ranges)
+        if self.renormalize:
+            outputs = (self._norm(output) for output in outputs)
+        return tuple(outputs)
+
+
 class MultiobjectiveModel(nn.Module):
     """Wrapper around contrastive learning models to all training with multiple objectives
 

cebra/solver/base.py

@@ -120,7 +120,7 @@ def load_state_dict(self, state_dict: dict, strict: bool = True):
                 to partially load the state for all given keys.
         """
 
-        def _contains(key):
+        def _contains(key, strict=strict):
             if key in state_dict:
                 return True
             elif strict:
@@ -146,7 +146,8 @@ def _get(key):
             self.decode_history = _get("decode")
         if _contains("log"):
             self.log = _get("log")
-        if _contains("metadata"):
+        # NOTE(stes): Added in CEBRA 0.6.0
+        if _contains("metadata", strict=False):
             self.metadata = _get("metadata")
 
     @property
@@ -405,7 +406,7 @@ def num_total_features(self):
     def __post_init__(self):
         super().__post_init__()
         self._check_dimensions()
-        self.model = cebra.models.MultiobjectiveModel(
+        self.model = cebra.models.LegacyMultiobjectiveModel(
             self.model,
             dimensions=(self.num_behavior_features, self.model.num_output),
             renormalize=self.renormalize_features,

tests/test_models.py

@@ -23,6 +23,7 @@
 
 import pytest
 import torch
+from torch import nn
 
 import cebra.models
 import cebra.models.model
@@ -87,6 +88,56 @@ def test_offset_models(model_name, batch_size, input_length):
         assert len(outputs) == batch_size
 
 
+def test_multiobjective():
+
+    # NOTE(stes): This test is deprecated and will be removed in a future version.
+    # As of CEBRA 0.6.0, the multi objective models are tested separately in
+    # test_multiobjective.py.
+
+    class TestModel(cebra.models.Model):
+
+        def __init__(self):
+            super().__init__(num_input=10, num_output=10)
+            self._model = nn.Linear(self.num_input, self.num_output)
+
+        def forward(self, x):
+            return self._model(x)
+
+        @property
+        def get_offset(self):
+            return None
+
+    model = TestModel()
+
+    multi_model_overlap = cebra.models.LegacyMultiobjectiveModel(
+        model,
+        dimensions=(4, 6),
+        output_mode="overlapping",
+        append_last_dimension=True)
+    multi_model_separate = cebra.models.LegacyMultiobjectiveModel(
+        model,
+        dimensions=(4, 6),
+        output_mode="separate",
+        append_last_dimension=True)
+
+    x = torch.randn(5, 10)
+
+    assert model(x).shape == (5, 10)
+
+    assert model.num_output == multi_model_overlap.num_output
+    assert model.get_offset == multi_model_overlap.get_offset
+
+    first, second, third = multi_model_overlap(x)
+    assert first.shape == (5, 4)
+    assert second.shape == (5, 6)
+    assert third.shape == (5, 10)
+
+    first, second, third = multi_model_separate(x)
+    assert first.shape == (5, 4)
+    assert second.shape == (5, 2)
+    assert third.shape == (5, 4)
+
+
 @pytest.mark.parametrize("version,raises", [
     ["1.12", False],
     ["2.", False],