Improve modularity remove duplicate code and todos

Author: CeliaBenquet

cebra/integrations/sklearn/cebra.py

@@ -1235,7 +1235,7 @@ def transform(self,
         sklearn_utils_validation.check_is_fitted(self, "n_features_")
         self.solver_._check_is_session_id_valid(session_id=session_id)
 
-        if torch.is_tensor(X):
+        if torch.is_tensor(X) and X.device.type == "cuda":
             X = X.detach().cpu()
 
         X = sklearn_utils.check_input_array(X, min_samples=len(self.offset_))
@@ -1256,61 +1256,6 @@ def transform(self,
 
         return output.detach().cpu().numpy()
 
-    # Deprecated, kept for testing.
-    def transform_deprecated(self,
-                             X: Union[npt.NDArray, torch.Tensor],
-                             session_id: Optional[int] = None) -> npt.NDArray:
-        """Transform an input sequence and return the embedding.
-
-        Args:
-            X: A numpy array or torch tensor of size ``time x dimension``.
-            session_id: The session ID, an :py:class:`int` between 0 and :py:attr:`num_sessions` for
-                multisession, set to ``None`` for single session.
-
-        Returns:
-            A :py:func:`numpy.array` of size ``time x output_dimension``.
-
-        Example:
-
-            >>> import cebra
-            >>> import numpy as np
-            >>> dataset =  np.random.uniform(0, 1, (1000, 30))
-            >>> cebra_model = cebra.CEBRA(max_iterations=10)
-            >>> cebra_model.fit(dataset)
-            CEBRA(max_iterations=10)
-            >>> embedding = cebra_model.transform(dataset)
-
-        """
-
-        sklearn_utils_validation.check_is_fitted(self, "n_features_")
-        model, offset = self._select_model(X, session_id)
-
-        # Input validation
-        X = sklearn_utils.check_input_array(X, min_samples=len(self.offset_))
-        #input_dtype = X.dtype
-
-        with torch.no_grad():
-            model.eval()
-
-            if self.pad_before_transform:
-                X = np.pad(X, ((offset.left, offset.right - 1), (0, 0)),
-                           mode="edge")
-            X = torch.from_numpy(X).float().to(self.device_)
-
-            if isinstance(model, cebra.models.ConvolutionalModelMixin):
-                # Fully convolutional evaluation, switch (T, C) -> (1, C, T)
-                X = X.transpose(1, 0).unsqueeze(0)
-                output = model(X).cpu().numpy().squeeze(0).transpose(1, 0)
-            else:
-                # Standard evaluation, (T, C, dt)
-                output = model(X).cpu().numpy()
-
-        #TODO: check if this is safe.
-        return output.numpy(force=True)
-
-        #if input_dtype == "float64":
-        #    return output.astype(input_dtype)
-
     def fit_transform(
         self,
         X: Union[npt.NDArray, torch.Tensor],

cebra/solver/base.py

@@ -451,6 +451,8 @@ def fit(
                 if logdir is not None:
                     self.save(logdir, f"checkpoint_{num_steps:#07d}.pth")
 
+        self._set_fitted_params(loader)
+
     def step(self, batch: cebra.data.Batch) -> dict:
         """Perform a single gradient update.
 
@@ -610,10 +612,6 @@ def transform(self,
 
         if len(offset) < 2 and pad_before_transform:
             pad_before_transform = False
-            
-        if batch_size is not None and batch_size < 1:
-            raise ValueError(
-                f"Batch size should be at least 1, got {batch_size}")
 
         model.eval()
         if batch_size is not None:
@@ -665,6 +663,12 @@ def load(self, logdir, filename="checkpoint.pth"):
         checkpoint = torch.load(savepath, map_location=self.device)
         self.load_state_dict(checkpoint, strict=True)
 
+        if hasattr(self.model, "n_features"):
+            n_features = self.model.n_features
+            self.n_features = ([
+                session_n_features for session_n_features in n_features
+            ] if isinstance(n_features, list) else n_features)
+
     def save(self, logdir, filename="checkpoint_last.pth"):
         """Save the model and optimizer params.
 

cebra/solver/multi_session.py

@@ -41,9 +41,9 @@ class MultiSessionSolver(abc_.Solver):
 
     def parameters(self, session_id: Optional[int] = None):
         """Iterate over all parameters."""
-        if session_id is not None:
-            for parameter in self.model[session_id].parameters():
-                yield parameter
+        self._check_is_session_id_valid(session_id=session_id)
+        for parameter in self.model[session_id].parameters():
+            yield parameter
 
         for parameter in self.criterion.parameters():
             yield parameter