Skip to content

extractor

shrinkai.adapters.extractor

Classes:

Name Description
FeatureExtractor

Wraps a PyTorch model to extract intermediate feature maps via forward hooks.

Classes

FeatureExtractor

Bases: Module

Wraps a PyTorch model to extract intermediate feature maps via forward hooks.

This wrapper does not modify the original model's source code. It dynamically attaches hooks to intercept the output of specified layers during the forward pass.

Methods:

Name Description
__init__

Initializes the FeatureExtractor.
forward

Performs a forward pass and captures intermediate features.
remove_hooks

Removes all registered hooks to prevent memory leaks.
Source code in src/shrinkai/adapters/extractor.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
class FeatureExtractor(nn.Module):
    """Wraps a PyTorch model to extract intermediate feature maps via forward hooks.

    This wrapper does not modify the original model's source code. It dynamically
    attaches hooks to intercept the output of specified layers during the forward pass.
    """

    def __init__(self, model: nn.Module, target_layers: Iterable[str] | Mapping[str, str]) -> None:
        """Initializes the FeatureExtractor.

        Args:
            model: The PyTorch model to extract features from.
            target_layers: An iterable of layer names to hook (e.g., ["layer1", "layer2"]),
                or a mapping dictionary to assign common aliases for cross-model matching
                (e.g., {"layer4.conv1": "block_1", "layer8.conv1": "block_2"}).

        Raises:
            ValueError: If a target layer does not exist in the model.
        """
        super().__init__()
        self.model = model
        if isinstance(target_layers, Mapping):
            aliases = list(target_layers.values())
            if len(set(aliases)) != len(aliases):
                raise ValueError("Many layers have the same alias.")
            self.target_layers = target_layers
        else:
            self.target_layers = {layer: layer for layer in target_layers}
        self.features: dict[str, torch.Tensor] = {}
        self._hooks: list[torch.utils.hooks.RemovableHandle] = []

        self._register_hooks()

    def _register_hooks(self) -> None:
        """Registers PyTorch forward hooks on the target layers."""
        model_modules = dict(self.model.named_modules())

        for layer_name, alias in self.target_layers.items():
            if layer_name not in model_modules:
                raise ValueError(
                    f"Layer '{layer_name}' not found in the model. "
                    f"Available layers: {list(model_modules.keys())[:10]}..."
                )

            module = model_modules[layer_name]

            def hook_fn(
                module: nn.Module, input: tuple, output: torch.Tensor, name: str = alias
            ) -> None:
                if isinstance(output, tuple):
                    output = output[0]  # take main tensor in first position
                self.features[name] = output

            handle = module.register_forward_hook(hook_fn)
            self._hooks.append(handle)

    def forward(
        self, x: torch.Tensor, *args: Any, **kwargs: Any
    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
        """Performs a forward pass and captures intermediate features.

        Args:
            x: Input tensor.
            *args: Additional positional arguments for the model.
            **kwargs: Additional keyword arguments for the model.

        Returns:
            tuple[torch.Tensor, dict[str, torch.Tensor]]: A tuple containing the
                final model output (logits) and a dictionary of extracted features.
        """
        self.features.clear()
        logits = self.model(x, *args, **kwargs)
        return logits, self.features.copy()

    def remove_hooks(self) -> None:
        """Removes all registered hooks to prevent memory leaks."""
        for handle in self._hooks:
            handle.remove()
        self._hooks.clear()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.remove_hooks()
Methods:
__init__ 
__init__(
    model: Module,
    target_layers: Iterable[str] | Mapping[str, str],
) -> None

Initializes the FeatureExtractor.

Parameters:

Name Type Description Default
model Module

The PyTorch model to extract features from.

 required
target_layers Iterable[str] | Mapping[str, str]

An iterable of layer names to hook (e.g., ["layer1", "layer2"]), or a mapping dictionary to assign common aliases for cross-model matching (e.g., {"layer4.conv1": "block_1", "layer8.conv1": "block_2"}).

 required

Raises:

Type Description
ValueError

If a target layer does not exist in the model.
Source code in src/shrinkai/adapters/extractor.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
def __init__(self, model: nn.Module, target_layers: Iterable[str] | Mapping[str, str]) -> None:
    """Initializes the FeatureExtractor.

    Args:
        model: The PyTorch model to extract features from.
        target_layers: An iterable of layer names to hook (e.g., ["layer1", "layer2"]),
            or a mapping dictionary to assign common aliases for cross-model matching
            (e.g., {"layer4.conv1": "block_1", "layer8.conv1": "block_2"}).

    Raises:
        ValueError: If a target layer does not exist in the model.
    """
    super().__init__()
    self.model = model
    if isinstance(target_layers, Mapping):
        aliases = list(target_layers.values())
        if len(set(aliases)) != len(aliases):
            raise ValueError("Many layers have the same alias.")
        self.target_layers = target_layers
    else:
        self.target_layers = {layer: layer for layer in target_layers}
    self.features: dict[str, torch.Tensor] = {}
    self._hooks: list[torch.utils.hooks.RemovableHandle] = []

    self._register_hooks()
forward 
forward(
    x: Tensor, *args: Any, **kwargs: Any
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]

Performs a forward pass and captures intermediate features.

Parameters:

Name Type Description Default
x Tensor

Input tensor.

 required
*args Any

Additional positional arguments for the model.

 ()
**kwargs Any

Additional keyword arguments for the model.

 {}

Returns:

Type Description
tuple[Tensor, dict[str, Tensor]]

tuple[torch.Tensor, dict[str, torch.Tensor]]: A tuple containing the final model output (logits) and a dictionary of extracted features.
Source code in src/shrinkai/adapters/extractor.py 
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
def forward(
    self, x: torch.Tensor, *args: Any, **kwargs: Any
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
    """Performs a forward pass and captures intermediate features.

    Args:
        x: Input tensor.
        *args: Additional positional arguments for the model.
        **kwargs: Additional keyword arguments for the model.

    Returns:
        tuple[torch.Tensor, dict[str, torch.Tensor]]: A tuple containing the
            final model output (logits) and a dictionary of extracted features.
    """
    self.features.clear()
    logits = self.model(x, *args, **kwargs)
    return logits, self.features.copy()
remove_hooks 
remove_hooks() -> None

Removes all registered hooks to prevent memory leaks.

Source code in src/shrinkai/adapters/extractor.py 
82
83
84
85
86
def remove_hooks(self) -> None:
    """Removes all registered hooks to prevent memory leaks."""
    for handle in self._hooks:
        handle.remove()
    self._hooks.clear()