Fix: Isolate hooks between multiple TorchFeatureExtractor instances

oodeel/extractor/hf_torch_feature_extractor.py

@@ -152,7 +152,10 @@ def predict_tensor(
         """
         if x.device != self._device:
             x = x.to(self._device)
+
+        self._active = True  # Activate hooks
         outputs = self.model(x, output_hidden_states=True, return_dict=True)
+        self._active = False  # Deactivate hooks
 
         features = []
         for feature_layer_id in self.feature_layers_id:

oodeel/extractor/torch_feature_extractor.py

@@ -88,6 +88,8 @@ def __init__(
         if return_penultimate:
             feature_layers_id.append("penultimate")
 
+        self._handles = []
+
         super().__init__(
             model=model,
             feature_layers_id=feature_layers_id,
@@ -101,6 +103,7 @@ def __init__(
         self._device = next(model.parameters()).device
         self._features = {layer: torch.empty(0) for layer in self._hook_layers_id}
         self._last_logits = None
+        self._active = False
         self.backend = "torch"
 
     @property
@@ -120,6 +123,8 @@ def _get_features_hook(self, layer_id: Union[str, int]) -> Callable:
         """
 
         def hook(_, __, output):
+            if not self._active:
+                return
             if isinstance(output, torch.Tensor):
                 self._features[layer_id] = output
             else:
@@ -140,6 +145,8 @@ def _get_penultimate_hook(self) -> Callable:
         """
 
         def hook(_, input):
+            if not self._active:
+                return
             if isinstance(input[0], torch.Tensor):
                 self._features["penultimate"] = input[0]
             else:
@@ -185,14 +192,10 @@ def find_layer(
 
     def prepare_extractor(self) -> None:
         """Prepare the feature extractor by adding hooks to self.model"""
-        # prepare self.model for ood hooks (add _ood_handles attribute or
-        # remove ood forward hooks attached to the model)
-        self._prepare_ood_handles()
-
         # === If react method, clip activations from penultimate layer ===
         if self.react_threshold is not None:
             pen_layer = self.find_layer(self.model, self.head_layer_id)
-            self.model._ood_handles.append(
+            self._handles.append(
                 pen_layer.register_forward_pre_hook(
                     self._get_clip_hook(self.react_threshold)
                 )
@@ -201,7 +204,7 @@ def prepare_extractor(self) -> None:
         # === If SCALE method, scale activations from penultimate layer ===
         if self.scale_percentile is not None:
             pen_layer = self.find_layer(self.model, self.head_layer_id)
-            self.model._ood_handles.append(
+            self._handles.append(
                 pen_layer.register_forward_pre_hook(
                     self._get_scale_hook(self.scale_percentile)
                 )
@@ -210,7 +213,7 @@ def prepare_extractor(self) -> None:
         # === If ASH method, scale and prune activations from penultimate layer ===
         if self.ash_percentile is not None:
             pen_layer = self.find_layer(self.model, self.head_layer_id)
-            self.model._ood_handles.append(
+            self._handles.append(
                 pen_layer.register_forward_pre_hook(
                     self._get_ash_hook(self.ash_percentile)
                 )
@@ -221,13 +224,13 @@ def prepare_extractor(self) -> None:
             if layer_id == "penultimate":
                 # Register penultimate hook
                 layer = self.find_layer(self.model, self.head_layer_id)
-                self.model._ood_handles.append(
+                self._handles.append(
                     layer.register_forward_pre_hook(self._get_penultimate_hook())
                 )
                 continue
 
             layer = self.find_layer(self.model, layer_id)
-            self.model._ood_handles.append(
+            self._handles.append(
                 layer.register_forward_hook(self._get_features_hook(layer_id))
             )
 
@@ -279,8 +282,10 @@ def predict_tensor(
         if x.device != self._device:
             x = x.to(self._device)
 
+        self._active = True
         with torch.set_grad_enabled(not detach):
             _ = self.model(x)
+        self._active = False
 
         if detach:
             features = [
@@ -415,6 +420,8 @@ def _get_clip_hook(self, threshold: float) -> Callable:
         """
 
         def hook(_, input):
+            if not self._active:
+                return
             input = input[0]
             input = torch.clip(input, max=threshold)
             return input
@@ -433,6 +440,8 @@ def _get_scale_hook(self, percentile: float) -> Callable:
         """
 
         def hook(_, input):
+            if not self._active:
+                return
             input = input[0]
             output_percentile = torch.quantile(input, percentile, dim=1)
             mask = input > output_percentile[:, None]
@@ -456,6 +465,8 @@ def _get_ash_hook(self, percentile: float) -> Callable:
         """
 
         def hook(_, input):
+            if not self._active:
+                return
             input = input[0]
             output_percentile = torch.quantile(input, percentile, dim=1)
             mask = input > output_percentile[:, None]
@@ -467,19 +478,11 @@ def hook(_, input):
 
         return hook
 
-    def _prepare_ood_handles(self) -> None:
-        """
-        Prepare the model by either setting a new attribute to self.model
-        as a list which will contain all the ood specific hooks, or by cleaning
-        the existing ood specific hooks if the attribute already exists.
-        """
-
-        if not hasattr(self.model, "_ood_handles"):
-            setattr(self.model, "_ood_handles", [])
-        else:
-            for handle in self.model._ood_handles:
-                handle.remove()
-            self.model._ood_handles = []
+    def clean_hooks(self) -> None:
+        """Remove hooks registered by this extractor instance from the model."""
+        for handle in self._handles:
+            handle.remove()
+        self._handles = []
 
     def _default_postproc_fn(self, feat: TensorType) -> TensorType:
         """Default postprocessing function to apply to each feature immediately

tests/tests_torch/extractor/test_torch_feature_extractor.py

@@ -21,6 +21,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 import pytest
+import torch
 import torch.nn as nn
 from torch.utils.data import DataLoader
 
@@ -199,3 +200,110 @@ def globalavg(x):
     feat0, feat1 = feats
     assert list(feat0.size()) == [100, 16, 1, 1]
     assert list(feat1.size()) == [100, 84]
+
+
+def test_multiple_extractors_hook_isolation():
+    """Test that multiple extractors on the same model don't interfere with each other.
+
+    This test verifies that when two TorchFeatureExtractor instances are created
+    on the same model, each extractor only captures features from its own hooks
+    when predict is called, without being affected by the other extractor's hooks.
+    """
+    n_samples = 20
+    input_shape = (3, 32, 32)
+    num_labels = 10
+
+    x = generate_data_torch(input_shape, num_labels, n_samples)
+    dataset = DataLoader(x, batch_size=n_samples)
+
+    model = Net()
+
+    # Create first extractor targeting fc1
+    extractor1 = TorchFeatureExtractor(model, feature_layers_id=["fc1"])
+
+    # Create second extractor targeting fc2
+    extractor2 = TorchFeatureExtractor(model, feature_layers_id=["fc2"])
+
+    # Get features from extractor1
+    features1, _ = extractor1.predict(dataset)
+
+    # Get features from extractor2
+    features2, _ = extractor2.predict(dataset)
+
+    # Verify each extractor captured the correct layer dimensions
+    # fc1 outputs 120 features, fc2 outputs 84 features
+    assert features1[0].shape == (n_samples, 120), "Extractor1 should capture fc1 (120)"
+    assert features2[0].shape == (n_samples, 84), "Extractor2 should capture fc2 (84)"
+
+    # Verify extractor1 still works after extractor2 was used
+    features1_again, _ = extractor1.predict(dataset)
+    assert torch.allclose(
+        features1[0], features1_again[0]
+    ), "Extractor1 should produce same results after extractor2 was used"
+
+
+def test_extractor_clean_hooks():
+    """Test that clean_hooks properly removes hooks from the model."""
+    n_samples = 10
+    input_shape = (3, 32, 32)
+    num_labels = 10
+
+    x = generate_data_torch(input_shape, num_labels, n_samples)
+    dataset = DataLoader(x, batch_size=n_samples)
+
+    model = Net()
+
+    # Create extractor and verify it works
+    extractor = TorchFeatureExtractor(model, feature_layers_id=["fc1"])
+    features, _ = extractor.predict(dataset)
+    assert features[0].shape == (n_samples, 120)
+
+    # Count hooks before cleanup
+    num_handles_before = len(extractor._handles)
+    assert num_handles_before > 0, "Extractor should have registered hooks"
+
+    # Clean hooks
+    extractor.clean_hooks()
+
+    # Verify that model's hooks are removed
+    for module in model.modules():
+        assert not module._forward_hooks, "All forward hooks should be removed"
+        assert not module._backward_hooks, "All backward hooks should be removed"
+        assert not module._forward_pre_hooks, "All forward pre-hooks should be removed"
+
+    # Verify hooks are removed
+    assert len(extractor._handles) == 0, "Hooks should be removed after clean_hooks"
+
+
+def test_react_hook_isolation():
+    """Test that ReAct hooks are also isolated between extractors."""
+    n_samples = 10
+    input_shape = (3, 32, 32)
+    num_labels = 10
+
+    x = generate_data_torch(input_shape, num_labels, n_samples)
+    dataset = DataLoader(x, batch_size=n_samples)
+
+    model = Net()
+
+    # Create extractor with ReAct (clips activations)
+    extractor_react = TorchFeatureExtractor(
+        model, feature_layers_id=["fc2"], react_threshold=1.0
+    )
+
+    # Create normal extractor
+    extractor_normal = TorchFeatureExtractor(model, feature_layers_id=["fc2"])
+
+    # Get features from normal extractor
+    features_normal, _ = extractor_normal.predict(dataset)
+
+    # Get features from react extractor
+    features_react, _ = extractor_react.predict(dataset)
+
+    # Get features from normal extractor again - should be unaffected by react
+    features_normal_again, _ = extractor_normal.predict(dataset)
+
+    # Normal extractor should produce same results before and after react was used
+    assert torch.allclose(
+        features_normal[0], features_normal_again[0]
+    ), "Normal extractor should not be affected by ReAct extractor"