diff --git a/oodeel/extractor/hf_torch_feature_extractor.py b/oodeel/extractor/hf_torch_feature_extractor.py index b6c75d98..2b20018d 100644 --- a/oodeel/extractor/hf_torch_feature_extractor.py +++ b/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: diff --git a/oodeel/extractor/torch_feature_extractor.py b/oodeel/extractor/torch_feature_extractor.py index 7bcbb399..cb1c3dc2 100644 --- a/oodeel/extractor/torch_feature_extractor.py +++ b/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 diff --git a/tests/tests_torch/extractor/test_torch_feature_extractor.py b/tests/tests_torch/extractor/test_torch_feature_extractor.py index 4d2ffbcf..40f5533d 100644 --- a/tests/tests_torch/extractor/test_torch_feature_extractor.py +++ b/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"