Comparator Extractor

extractor

CLASS	DESCRIPTION
FeatureExtractor	Extract intermediate features from neural networks using forward hooks.

FUNCTION	DESCRIPTION
extract_features	Extract features from specified layers of a model.

Classes

FeatureExtractor()

Extract intermediate features from neural networks using forward hooks.

This class manages PyTorch forward hooks to capture layer features during a forward pass through a neural network. Hooks are automatically cleaned up to prevent memory leaks.

ATTRIBUTE	DESCRIPTION
features	Dictionary mapping layer names to their feature tensors. TYPE: dict[str, Tensor]
hooks	List of registered hook handles. TYPE: list

Examples:

Basic usage with sequential model:

>>> import torch
>>> import torch.nn as nn
>>> model = nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.Linear(20, 5))
>>> X = torch.randn(100, 10)
>>> extractor = FeatureExtractor()
>>> extractor.register_hooks(model, ["0", "2"])
>>> model(X)
>>> features = extractor.get_features()
>>> extractor.clear()

Using with custom models:

>>> class Net(nn.Module):
...     def __init__(self):
...         super().__init__()
...         self.conv1 = nn.Conv2d(3, 64, 3)
...         self.fc = nn.Linear(64, 10)
...
...     def forward(self, x):
...         x = self.conv1(x)
...         return self.fc(x.flatten(1))
>>> model = Net()
>>> extractor = FeatureExtractor()
>>> extractor.register_hooks(model, ["conv1", "fc"])

Notes

Call clear() after extracting features to remove hooks.

METHOD	DESCRIPTION
register_hooks	Register forward hooks on specified layers.
get_features	Get extracted features.
clear	Remove all registered hooks and clear features.

Source code in spectre/comparator/extractor.py

def __init__(self):
    self.features: dict[str, torch.Tensor] = {}
    self.hooks: list = []

Functions

register_hooks(model: torch.nn.Module, layer_names: list[str]) -> None

Register forward hooks on specified layers.

PARAMETER	DESCRIPTION
model	Model to extract features from. TYPE: Module
layer_names	Names of layers to extract features from. TYPE: list[str]

RAISES	DESCRIPTION
ValueError	If a specified layer name is not found in the model.

Source code in spectre/comparator/extractor.py

def register_hooks(self, model: torch.nn.Module, layer_names: list[str]) -> None:
    """
    Register forward hooks on specified layers.

    Parameters
    ----------
    model : nn.Module
        Model to extract features from.

    layer_names : list[str]
        Names of layers to extract features from.

    Raises
    ------
    ValueError
        If a specified layer name is not found in the model.
    """
    module_dict = dict(model.named_modules())

    for name in layer_names:
        if name not in module_dict:
            raise ValueError(f"Layer '{name}' not found in model.")

        module = module_dict[name]
        hook = module.register_forward_hook(self._hook_fn(name))
        self.hooks.append(hook)

get_features() -> dict[str, torch.Tensor]

Get extracted features.

RETURNS	DESCRIPTION
dict[str, Tensor]	Dictionary mapping layer names to feature tensors.

Source code in spectre/comparator/extractor.py

def get_features(self) -> dict[str, torch.Tensor]:
    """
    Get extracted features.

    Returns
    -------
    dict[str, torch.Tensor]
        Dictionary mapping layer names to feature tensors.
    """
    return self.features.copy()

clear() -> None

Remove all registered hooks and clear features.

Source code in spectre/comparator/extractor.py

def clear(self) -> None:
    """Remove all registered hooks and clear features."""
    for hook in self.hooks:
        hook.remove()
    self.hooks.clear()
    self.features.clear()

Functions

extract_features(model: torch.nn.Module, X: torch.Tensor, layers: list[str] | None = None, device: torch.device | str | None = None, flatten_conv: bool = True) -> dict[str, torch.Tensor]

Extract features from specified layers of a model.

This function uses PyTorch forward hooks to capture intermediate features during a forward pass. Hooks are automatically cleaned up after extraction.

PARAMETER	DESCRIPTION
model	Model to extract features from. Can be a torch.nn.Module or Parametric model (will use .model attribute). TYPE: Module
X	Input data, shape (n_samples, ...). Must be compatible with model input. TYPE: Tensor
layers	Names of layers to extract features from. If None, extract from all leaf modules (Conv2d, Linear, ReLU, etc.). Layer names follow PyTorch's named_modules() convention. TYPE: list[str] | None, optional, by default None DEFAULT: None
device	Device to use for computation. If None, use model's device. TYPE: torch.device | str | None, optional, by default None DEFAULT: None
flatten_conv	If True, flatten spatial dimensions of convolutional layer outputs to shape (n_samples, channels * height * width). Useful for comparing convolutional and fully-connected representations. TYPE: bool, optional, by default True DEFAULT: True

RETURNS	DESCRIPTION
dict[str, Tensor]	Dictionary mapping layer names to feature tensors. Each feature has shape (n_samples, ...) depending on layer output. Features are detached from computation graph.

RAISES	DESCRIPTION
ValueError	If no layers are found to extract from, or if specified layer names don't exist in the model.

Examples:

Extract from all layers:

>>> import torch
>>> import torch.nn as nn
>>> model = nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.Linear(20, 5))
>>> X = torch.randn(100, 10)
>>> features = extract_features(model, X)
>>> features.keys()
dict_keys(['0', '1', '2'])
>>> features["0"].shape
torch.Size([100, 20])

Extract from specific layers:

>>> features = extract_features(model, X, layers=["0", "2"])
>>> features.keys()
dict_keys(['0', '2'])

With convolutional layers:

>>> conv_model = nn.Sequential(
...     nn.Conv2d(3, 64, 3), nn.ReLU(), nn.Flatten(), nn.Linear(64 * 26 * 26, 10)
... )
>>> X_img = torch.randn(32, 3, 28, 28)
>>> features = extract_features(conv_model, X_img, flatten_conv=True)
>>> features["0"].shape  # Flattened conv output
torch.Size([32, 43264])

Notes

Temporarily sets the model to eval mode during extraction and restores the original training state afterward. All features are detached from the computation graph to prevent gradient tracking.

Source code in spectre/comparator/extractor.py

def extract_features(
    model: torch.nn.Module,
    X: torch.Tensor,
    layers: list[str] | None = None,
    device: torch.device | str | None = None,
    flatten_conv: bool = True,
) -> dict[str, torch.Tensor]:
    """
    Extract features from specified layers of a model.

    This function uses PyTorch forward hooks to capture intermediate
    features during a forward pass. Hooks are automatically cleaned up
    after extraction.

    Parameters
    ----------
    model : nn.Module
        Model to extract features from. Can be a `torch.nn.Module` or
        `Parametric` model (will use `.model` attribute).

    X : torch.Tensor
        Input data, shape (n_samples, ...). Must be compatible with model input.

    layers : list[str] | None, optional, by default None
        Names of layers to extract features from. If None, extract from
        all leaf modules (Conv2d, Linear, ReLU, etc.). Layer names follow
        PyTorch's `named_modules()` convention.

    device : torch.device | str | None, optional, by default None
        Device to use for computation. If None, use model's device.

    flatten_conv : bool, optional, by default True
        If True, flatten spatial dimensions of convolutional layer outputs
        to shape (n_samples, channels * height * width). Useful for comparing
        convolutional and fully-connected representations.

    Returns
    -------
    dict[str, torch.Tensor]
        Dictionary mapping layer names to feature tensors.
        Each feature has shape (n_samples, ...) depending on layer output.
        Features are detached from computation graph.

    Raises
    ------
    ValueError
        If no layers are found to extract from, or if specified layer names
        don't exist in the model.

    Examples
    --------
    Extract from all layers:

    >>> import torch
    >>> import torch.nn as nn
    >>> model = nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.Linear(20, 5))
    >>> X = torch.randn(100, 10)
    >>> features = extract_features(model, X)
    >>> features.keys()
    dict_keys(['0', '1', '2'])
    >>> features["0"].shape
    torch.Size([100, 20])

    Extract from specific layers:

    >>> features = extract_features(model, X, layers=["0", "2"])
    >>> features.keys()
    dict_keys(['0', '2'])

    With convolutional layers:

    >>> conv_model = nn.Sequential(
    ...     nn.Conv2d(3, 64, 3), nn.ReLU(), nn.Flatten(), nn.Linear(64 * 26 * 26, 10)
    ... )
    >>> X_img = torch.randn(32, 3, 28, 28)
    >>> features = extract_features(conv_model, X_img, flatten_conv=True)
    >>> features["0"].shape  # Flattened conv output
    torch.Size([32, 43264])

    Notes
    -----
    Temporarily sets the model to eval mode during extraction and restores the original
    training state afterward. All features are detached from the computation graph
    to prevent gradient tracking.
    """
    # Validate input tensor
    if not isinstance(X, torch.Tensor):
        raise TypeError(f"X must be a torch.Tensor, got {type(X)}")
    if X.ndim < 2:
        raise ValueError(
            f"X must have at least 2 dimensions (n_samples, ...), got shape {X.shape}"
        )
    # Handle Parametric models
    if hasattr(model, "model") and isinstance(getattr(model, "model"), torch.nn.Module):
        model = model.model

    if device is None:
        # Try to get device from model parameters
        try:
            device = next(model.parameters()).device
        except StopIteration:
            # Model has no parameters, try buffers (e.g., BatchNorm running stats)
            try:
                device = next(model.buffers()).device
            except StopIteration:
                # No parameters or buffers, default to CPU
                device = torch.device("cpu")
    X = X.to(device)

    # Determine layers to extract
    if layers is None:
        # Extract from all leaf modules
        layers = []
        for name, module in model.named_modules():
            if name != "" and len(list(module.children())) == 0:
                layers.append(name)

        # Handle case where model itself is a leaf module (e.g., bare torch.nn.Linear)
        if len(layers) == 0:
            # Check if the model itself is a leaf module with parameters
            # (exclude empty containers like nn.Sequential())
            has_children = len(list(model.children())) > 0
            has_params = len(list(model.parameters())) > 0

            if not has_children and has_params:
                # Use empty string as name for the root module
                layers = [""]

    if len(layers) == 0:
        raise ValueError("No layers found to extract features from.")

    # Temporarily set model to eval mode to register hooks
    was_training = model.training
    model.eval()

    extractor = FeatureExtractor()
    try:
        extractor.register_hooks(model, layers)

        with torch.no_grad():
            _ = model(X)

        features = extractor.get_features()
    finally:
        extractor.clear()
        # Set model back to original training mode
        if was_training:
            model.train()

    # Optionally flatten convolutional features
    if flatten_conv:
        for name, act in features.items():
            if act.ndim > 2:
                # Flatten dim: (n_samples, C, H, W) -> (n_samples, C * H * W)
                features[name] = act.reshape(act.shape[0], -1)

    return features