Fix tests

* bump version

* update dockerfile

* fix progress bar

* remove outdated test

* rename models

Author: stes

Dockerfile

@@ -40,10 +40,10 @@ RUN make dist
 FROM cebra-base
 
 # install the cebra wheel
-ENV WHEEL=cebra-0.4.0-py2.py3-none-any.whl
+ENV WHEEL=cebra-0.4.0+regcl-py2.py3-none-any.whl
 WORKDIR /build
 COPY --from=wheel /build/dist/${WHEEL} .
-RUN pip install --no-cache-dir ${WHEEL}'[dev,integrations,datasets]'
+RUN pip install --no-cache-dir ${WHEEL}'[dev,integrations,datasets,regcl]'
 RUN rm -rf /build
 
 # add the repository

Makefile

@@ -1,4 +1,4 @@
-CEBRA_VERSION := 0.4.0
+CEBRA_VERSION := 0.4.0+regcl
 
 dist:
 	python3 -m pip install virtualenv

PKGBUILD

@@ -1,7 +1,7 @@
 # Maintainer: Steffen Schneider <[email protected]>
 pkgname=python-cebra
 _pkgname=cebra
-pkgver=0.4.0
+pkgver=0.4.0+regcl
 pkgrel=1
 pkgdesc="Consistent Embeddings of high-dimensional Recordings using Auxiliary variables"
 url="https://cebra.ai"

cebra/models/model.py

@@ -871,7 +871,7 @@ def get_offset(self) -> cebra.data.datatypes.Offset:
 
 
 @register("offset1-model-mse-tanh")
-class Offset0ModelMSE(_OffsetModel):
+class Offset0ModelMSETanH(_OffsetModel):
     """CEBRA model with a single sample receptive field, without output normalization."""
 
     def __init__(self, num_neurons, num_units, num_output, normalize=False):
@@ -901,7 +901,7 @@ def get_offset(self) -> cebra.data.datatypes.Offset:
 @parametrize("offset1-model-mse-clip-{clip_min}-{clip_max}",
              clip_min=(1000, 100, 50, 25, 20, 15, 10, 5, 1),
              clip_max=(1000, 100, 50, 25, 20, 15, 10, 5, 1))
-class Offset0ModelMSE(_OffsetModel):
+class Offset0ModelMSEClip(_OffsetModel):
     """CEBRA model with a single sample receptive field, without output normalization."""
 
     def __init__(self,
@@ -942,7 +942,7 @@ def get_offset(self) -> cebra.data.datatypes.Offset:
 @parametrize("offset1-model-mse-v2-{n_intermediate_layers}layers{tanh}",
              n_intermediate_layers=(1, 2, 3, 4, 5, 6, 7, 8, 9, 10),
              tanh=("-tanh", ""))
-class Offset0Model(_OffsetModel):
+class Offset0ModelMSETanHv2(_OffsetModel):
     """CEBRA model with a single sample receptive field, without output normalization."""
 
     def __init__(self,
@@ -993,7 +993,7 @@ def get_offset(self) -> cebra.data.datatypes.Offset:
 @parametrize("offset1-model-mse-resnet-{n_intermediate_layers}layers{tanh}",
              n_intermediate_layers=(1, 2, 3, 4, 5, 6, 7, 8, 9, 10),
              tanh=("-tanh", ""))
-class Offset0Model(_OffsetModel):
+class Offset0ModelResNetTanH(_OffsetModel):
     """CEBRA model with a single sample receptive field, without output normalization."""
 
     def __init__(self,

cebra/solver/util.py

@@ -90,15 +90,16 @@ def __post_init__(self):
             raise ValueError(
                 f"log_format must be one of {self._valid_formats}, "
                 f"but got {self.log_formats}")
+        self._stats = None
 
     def __iter__(self):
         self.iterator = self.loader
         if self.use_tqdm:
             self.iterator = tqdm.tqdm(self.iterator)
         for num_batch, batch in enumerate(self.iterator):
             yield num_batch, batch
-            self._log_message(num_batch, self.iterator.stats)
-        self._log_message(num_batch, self.iterator.stats)
+            self._log_message(num_batch, self._stats)
+        self._log_message(num_batch, self._stats)
 
     def _log_message(self, num_steps, stats):
         if self.logger is None:
@@ -119,4 +120,4 @@ def set_description(self, stats: Dict[str, float]):
         if self.use_tqdm:
             self.iterator.set_description(_description(stats))
 
-        self.iterator.stats = stats
+        self._stats = stats

reinstall.sh

@@ -15,7 +15,7 @@ pip uninstall -y cebra
 # Get version info after uninstalling --- this will automatically get the
 # most recent version based on the source code in the current directory.
 # $(tools/get_cebra_version.sh)
-VERSION=0.4.0
+VERSION=0.4.0+regcl
 echo "Upgrading to CEBRA v${VERSION}"
 
 # Upgrade the build system (PEP517/518 compatible)

tests/test_models.py

@@ -88,52 +88,6 @@ def test_offset_models(model_name, batch_size, input_length):
         assert len(outputs) == batch_size
 
 
-def test_multiobjective():
-
-    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.MultiobjectiveModel(
-        model,
-        dimensions=(4, 6),
-        output_mode="overlapping",
-        append_last_dimension=True)
-    multi_model_separate = cebra.models.MultiobjectiveModel(
-        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],