Adapt so that runs for unified CEBRA and simplification

Author: CeliaBenquet

cebra_lens/__init__.py

@@ -9,6 +9,7 @@
 from .utils_allen import *
 from .utils_hpc import *
 from .utils_plot import *
+from .utils_wrapper import *
 
 # selects what files can be imported when doing from CEBRA_Lens import * --> keep env clean
 # __all__ = ['get_layer_activations']

cebra_lens/activations.py

@@ -9,6 +9,8 @@
 import torch
 import torch.nn as nn
 
+from cebra_lens import utils_wrapper
+
 from .utils_plot import plot_activations
 
 
@@ -83,7 +85,8 @@ def get_cut_indices(
     elif layer_type == None:
         raise NotImplementedError(
             "Padding handling not implemented to handle activations for all layer types.",
-            "Set layer_type to nn.Conv1d to use the default padding handling.")
+            "Set layer_type to nn.Conv1d to use the default padding handling.",
+        )
     else:
         # need to analyze the padding from the last output of Conv1 and apply the same cut
         raise NotImplementedError(
@@ -94,8 +97,10 @@ def get_cut_indices(
 def get_activations_model(
     model: cebra.integrations.sklearn.cebra.CEBRA,
     data: torch.Tensor,
-    session_id: int = -1,
-    name: str = "single",
+    labels: Optional[torch.Tensor] = None,
+    session_id: int = None,
+    pad_before_transform: bool = True,
+    activations_keys_prefix: str = "model",
     instance: int = 0,
     layer_type: Type[nn.Module] = nn.Conv1d,
 ) -> Dict[str, npt.NDArray]:
@@ -129,42 +134,41 @@ def get_activations_model(
 
     activations = {}
     transform_kwargs = {}
-    if model.solver_name_ in [
-            "multi-session",
-            "multi-session-aux",
-            "multiobjective-solver",
-    ]:
-
-        model_ = model.model_[session_id]
-        transform_kwargs.update({"session_id": session_id})
-
-    elif model.solver_name_ in [
-            "single-session",
-            "single-session-aux",
-            "single-session-hybrid",
-            "single-session-full",
-    ]:
-        model_ = model.model_
 
+    if isinstance(model, cebra.integrations.sklearn.cebra.CEBRA):
+        model_ = model.solver_._get_model(session_id=session_id)
+    elif isinstance(model, cebra.solver.Solver):
+        model_ = model._get_model(session_id=session_id)
     else:
-        raise NotImplementedError(
-            f"Solver {model.solver_name_} is not yet implemented.")
+        raise ValueError(
+            "Model must be an instance of cebra.integrations.sklearn.cebra.CEBRA "
+            f"or cebra.solver.Solver, got {type(model)} instead.", )
+
+    transform_kwargs.update({"session_id": session_id})
 
     activations, handles, conv_layer_info = _attach_hooks(
         activations=activations,
         model=model_,
-        name=name,
+        activations_keys_prefix=activations_keys_prefix,
         instance=instance,
         layer_type=layer_type,
     )
-    _ = model.transform(data, **transform_kwargs)
+
+    _ = utils_wrapper.transform(model=model,
+                                data=data,
+                                label=labels,
+                                **transform_kwargs)
 
     # remove all handles to avoid activation's problems
     for handle in handles:
         handle.remove()
 
-    if model.pad_before_transform:
-        # Padding logic: calculate the total reduction which happens based on the kernel size per layer, divide the reduction per layer into 2 parts
+    if hasattr(model, "pad_before_transform"):
+        pad_before_transform = model.pad_before_transform
+
+    if pad_before_transform:
+        # Padding logic: calculate the total reduction which happens based on the
+        # kernel size per layer, divide the reduction per layer into 2 parts
         cut_indices = get_cut_indices(model_, layer_type, conv_layer_info)
         for i, (key, value) in enumerate(activations.items()):
             activations[key] = _cut_array(value, cut_indices[i])
@@ -175,7 +179,9 @@ def get_activations_model(
 def process_activations(
     models: Dict[str, List[cebra.integrations.sklearn.cebra.CEBRA]],
     data: torch.Tensor,
-    session_id: int,
+    labels: Optional[torch.Tensor] = None,
+    session_id: int = None,
+    pad_before_transform: bool = True,
     activations: Dict[str, npt.NDArray] = {},
     layer_type: Type[nn.Module] = None,
 ) -> Dict[str, npt.NDArray]:
@@ -209,8 +215,10 @@ def process_activations(
                 get_activations_model(
                     model=model,
                     data=data,
+                    labels=labels,
                     session_id=session_id,
-                    name=model_name,
+                    pad_before_transform=pad_before_transform,
+                    activations_keys_prefix=model_name,
                     instance=i,
                     layer_type=layer_type,
                 ))
@@ -219,18 +227,37 @@ def process_activations(
 
 
 # Function to create a hook that stores the activations in the dictionary
-def _get_activation(name: str, activations: Dict):
+def _get_activation(activations_keys_prefix: str, activations: Dict):
+    """Creates a forward hook to capture activations from a model layer.
+    
+    This function returns a hook that captures the output of a model layer during the forward pass and stores it in a dictionary.
+    
+    Parameters:
+    -----------
+    activations_keys_prefix : str
+        The prefix to use for the activation key, corresopnding to the type of model (eg. "single", "multi").
+    activations : Dict
+        A dictionary to store the activations. The key will be the name of the layer, and the value will be the activations.
+    
+    Returns:
+    --------
+    hook : function
+        A forward hook function that captures the activations.
+    activations : Dict
+        The dictionary where the activations will be stored. The key is the name of the layer, and the value is the activations.
+    """
 
     def hook(model, input, output):
-        activations[name] = output.detach().squeeze().numpy()
+        activations[activations_keys_prefix] = output.detach().squeeze().numpy(
+        )
 
     return hook, activations
 
 
 def _attach_hooks(
     activations: Dict[str, npt.NDArray],
     model: cebra.integrations.sklearn.cebra.CEBRA,
-    name: str,
+    activations_keys_prefix: str,
     instance: int,
     layer_type: Type[nn.Module] = None,
 ) -> Dict[str, npt.NDArray]:  # only attaches hooks on convolutional layers
@@ -244,8 +271,10 @@ def _attach_hooks(
         A dictionary to store the activations. Please refer to ``activations`` returned by ``get_activations_model``.
     model : cebra.integrations.sklearn.cebra.CEBRA
         The model to which hooks will be attached.
-    name : str
-        A base name for the activation keys (e.g., "single", "multi").
+    activations_keys_prefix : str
+        A base name for the activation keys (e.g., "single", "multi") so that the keys are 
+        unique for each model instance. The keys will be in the format 
+        '{activations_keys_prefix}_{instance}_layer_{num_layer}'.
     instance : int
         The instance number for the model, used to differentiate between models from the same model category.
     layer_type : Type[nn.Module]
@@ -266,7 +295,8 @@ def _attach_hooks(
             # attach hook to the layer_type and to the output layer
             if isinstance(model.net[i], layer_type) or i == len(model.net) - 1:
                 hook, activations = _get_activation(
-                    f"{name}_{instance}_layer_{num_layer}", activations)
+                    f"{activations_keys_prefix}_{instance}_layer_{num_layer}",
+                    activations)
                 if isinstance(model.net[i], layer_type):
                     conv_layer_info.append(model.net[i].kernel_size[0])
                 handle = model.net[i].register_forward_hook(hook)
@@ -277,7 +307,7 @@ def _attach_hooks(
                 for submodule in model.net[i].modules():
                     if isinstance(submodule, layer_type):
                         hook, activations = _get_activation(
-                            f"{name}_{instance}_layer_{num_layer}",
+                            f"{activations_keys_prefix}_{instance}_layer_{num_layer}",
                             activations,
                         )
                         conv_layer_info.append(submodule.kernel_size[0])
@@ -292,7 +322,7 @@ def _attach_hooks(
             if bool(model.net[i]._modules):
                 for submodule in model.net[i].modules():
                     hook, activations = _get_activation(
-                        f"{name}_{instance}_layer_{num_layer}",
+                        f"{activations_keys_prefix}_{instance}_layer_{num_layer}",
                         activations,
                     )
                     handle = submodule.register_forward_hook(hook)
@@ -301,7 +331,8 @@ def _attach_hooks(
 
             else:
                 hook, activations = _get_activation(
-                    f"{name}_{instance}_layer_{num_layer}", activations)
+                    f"{activations_keys_prefix}_{instance}_layer_{num_layer}",
+                    activations)
 
                 handle = model.net[i].register_forward_hook(hook)
                 handles.append(handle)
@@ -356,9 +387,11 @@ def aggregate_activations(
 def get_activations(
     models: Dict[str, List[cebra.integrations.sklearn.cebra.CEBRA]],
     data: torch.Tensor,
-    session_id: int,
+    labels: Optional[torch.Tensor] = None,
+    session_id: int = None,
+    pad_before_transform: bool = True,
     activations: Optional[Dict[str, npt.NDArray]] = None,
-    layer_type: Optional[Type[nn.Module]] = None,
+    layer_type: Optional[Type[nn.Module]] = nn.Conv1d,
 ) -> Dict[str, npt.NDArray]:
     """
     Extracts and organizes activations from models.
@@ -388,7 +421,15 @@ def get_activations(
     activations = activations or {}
 
     aggregated_activations = aggregate_activations(
-        process_activations(models, data, session_id, activations, layer_type))
+        process_activations(
+            models=models,
+            data=data,
+            labels=labels,
+            session_id=session_id,
+            pad_before_transform=pad_before_transform,
+            activations=activations,
+            layer_type=layer_type,
+        ))
 
     activations_dict = {}
     for key, value in aggregated_activations.items():

cebra_lens/quantification/cka_metric.py

@@ -212,7 +212,7 @@ def compute(self, activations: Dict[str, npt.NDArray],
         Parameters:
         -----------
         activations : Dict[str, npt.NDArray]
-            A dictionary where keys are strings which represent the model label and values are 2d lists 
+            A dictionary where keys are strings which represent the model label and values are 2d lists
             with the corresponding activations per layer.
 
         comparison : Tuple[str, str]