oodeel 0.4.0__py3-none-any.whl

oodeel/extractor/torch_feature_extractor.py

@@ -0,0 +1,506 @@
+# -*- coding: utf-8 -*-
+# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+from typing import get_args
+from typing import Optional
+
+import numpy as np
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from ..datasets.torch_data_handler import TorchDataHandler
+from ..types import Callable
+from ..types import ItemType
+from ..types import List
+from ..types import TensorType
+from ..types import Tuple
+from ..types import Union
+from ..utils.torch_operator import sanitize_input
+from .feature_extractor import FeatureExtractor
+
+
+class TorchFeatureExtractor(FeatureExtractor):
+    """
+    Feature extractor based on "model" to construct a feature space
+    on which OOD detection is performed. The features can be the output
+    activation values of internal model layers,
+    or the output of the model (logits).
+
+    Args:
+        model: model to extract the features from
+        feature_layers_id: list of str or int that identify features to output.
+            If int, the rank of the layer in the layer list
+            If str, the name of the layer. Defaults to [].
+        head_layer_id (int, str): identifier of the head layer.
+            If int, the rank of the layer in the layer list
+            If str, the name of the layer.
+            Defaults to -1
+        input_layer_id: input layer of the feature extractor (to avoid useless forwards
+            when working on the feature space without finetuning the bottom of
+            the model).
+            Defaults to None.
+        react_threshold: if not None, penultimate layer activations are clipped under
+            this threshold value (useful for ReAct). Defaults to None.
+        scale_percentile: if not None, the features are scaled
+            following the method of Xu et al., ICLR 2024.
+            Defaults to None.
+        ash_percentile: if not None, the features are scaled following
+            the method of Djurisic et al., ICLR 2023.
+        return_penultimate (bool): if True, the penultimate values are returned,
+            i.e. the input to the head_layer.
+
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        feature_layers_id: List[Union[int, str]] = [],
+        head_layer_id: Optional[Union[int, str]] = -1,
+        input_layer_id: Optional[Union[int, str]] = None,
+        react_threshold: Optional[float] = None,
+        scale_percentile: Optional[float] = None,
+        ash_percentile: Optional[float] = None,
+        return_penultimate: Optional[bool] = False,
+    ):
+        model = model.eval()
+
+        if return_penultimate:
+            feature_layers_id.append("penultimate")
+
+        super().__init__(
+            model=model,
+            feature_layers_id=feature_layers_id,
+            head_layer_id=head_layer_id,
+            input_layer_id=input_layer_id,
+            react_threshold=react_threshold,
+            scale_percentile=scale_percentile,
+            ash_percentile=ash_percentile,
+            return_penultimate=return_penultimate,
+        )
+        self._device = next(model.parameters()).device
+        self._features = {layer: torch.empty(0) for layer in self._hook_layers_id}
+        self._last_logits = None
+        self.backend = "torch"
+
+    @property
+    def _hook_layers_id(self):
+        return self.feature_layers_id + [self.head_layer_id]
+
+    def _get_features_hook(self, layer_id: Union[str, int]) -> Callable:
+        """
+        Hook that stores features corresponding to a specific layer
+        in a class dictionary.
+
+        Args:
+            layer_id (Union[str, int]): layer identifier
+
+        Returns:
+            Callable: hook function
+        """
+
+        def hook(_, __, output):
+            if isinstance(output, torch.Tensor):
+                self._features[layer_id] = output
+            else:
+                raise NotImplementedError
+
+        return hook
+
+    def _get_penultimate_hook(self) -> Callable:
+        """
+        Hook that stores features corresponding to a specific layer
+        in a class dictionary.
+
+        Args:
+            layer_id (Union[str, int]): layer identifier
+
+        Returns:
+            Callable: hook function
+        """
+
+        def hook(_, input):
+            if isinstance(input[0], torch.Tensor):
+                self._features["penultimate"] = input[0]
+            else:
+                raise NotImplementedError
+
+        return hook
+
+    @staticmethod
+    def find_layer(
+        model: nn.Module,
+        layer_id: Union[str, int],
+        index_offset: int = 0,
+        return_id: bool = False,
+    ) -> Union[nn.Module, Tuple[nn.Module, str]]:
+        """Find a layer in a model either by his name or by his index.
+
+        Args:
+            model (nn.Module): model whose identified layer will be returned
+            layer_id (Union[str, int]): layer identifier
+            index_offset (int): index offset to find layers located before (negative
+                offset) or after (positive offset) the identified layer
+            return_id (bool): if True, the layer will be returned with its id
+
+        Returns:
+            Union[nn.Module, Tuple[nn.Module, str]]: the corresponding layer and its id
+                if return_id is True.
+        """
+        if isinstance(layer_id, int):
+            layer_id += index_offset
+            if isinstance(model, nn.Sequential):
+                layer = model[layer_id]
+            else:
+                layer = list(model.named_modules())[layer_id][1]
+        else:
+            layer_id = list(dict(model.named_modules()).keys()).index(layer_id)
+            layer_id += index_offset
+            layer = list(model.named_modules())[layer_id][1]
+
+        if return_id:
+            return layer, layer_id
+        else:
+            return 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(
+                pen_layer.register_forward_pre_hook(
+                    self._get_clip_hook(self.react_threshold)
+                )
+            )
+
+        # === 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(
+                pen_layer.register_forward_pre_hook(
+                    self._get_scale_hook(self.scale_percentile)
+                )
+            )
+
+        # === 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(
+                pen_layer.register_forward_pre_hook(
+                    self._get_ash_hook(self.ash_percentile)
+                )
+            )
+
+        # Register a hook to store feature values for each considered layer + last layer
+        for layer_id in self._hook_layers_id:
+            if layer_id == "penultimate":
+                # Register penultimate hook
+                layer = self.find_layer(self.model, self.head_layer_id)
+                self.model._ood_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(
+                layer.register_forward_hook(self._get_features_hook(layer_id))
+            )
+
+        # Crop model if input layer is provided
+        if not (self.input_layer_id) is None:
+            if isinstance(self.input_layer_id, int):
+                if isinstance(self.model, nn.Sequential):
+                    self.model = nn.Sequential(
+                        *list(self.model.modules())[self.input_layer_id :]
+                    )
+                else:
+                    raise NotImplementedError
+            elif isinstance(self.input_layer_id, str):
+                if isinstance(self.model, nn.Sequential):
+                    module_names = list(
+                        filter(
+                            lambda x: x != "",
+                            map(lambda x: x[0], self.model.named_modules()),
+                        )
+                    )
+                    input_module_idx = module_names.index(self.input_layer_id)
+                    self.model = nn.Sequential(
+                        *list(self.model.modules())[(input_module_idx + 1) :]
+                    )
+                else:
+                    raise NotImplementedError
+            else:
+                raise NotImplementedError
+
+    @sanitize_input
+    def predict_tensor(
+        self,
+        x: TensorType,
+        postproc_fns: Optional[List[Callable]] = None,
+        detach: bool = True,
+    ) -> Tuple[List[torch.Tensor], torch.Tensor]:
+        """Get the projection of tensor in the feature space of self.model
+
+        Args:
+            x (TensorType): input tensor (or dataset elem)
+            postproc_fns (Optional[List[Callable]]): postprocessing function to apply to
+                each feature immediately after forward. Default to None.
+            detach (bool): if True, return features detached from the computational
+                graph. Defaults to True.
+
+        Returns:
+            List[torch.Tensor], torch.Tensor: features, logits
+        """
+        if x.device != self._device:
+            x = x.to(self._device)
+
+        with torch.set_grad_enabled(not detach):
+            _ = self.model(x)
+
+        if detach:
+            features = [
+                self._features[layer_id].detach() for layer_id in self._hook_layers_id
+            ]
+        else:
+            features = [self._features[layer_id] for layer_id in self._hook_layers_id]
+
+        # split features and logits
+        logits = features.pop()
+
+        if postproc_fns is not None:
+            features = [
+                postproc_fn(feature)
+                for feature, postproc_fn in zip(features, postproc_fns)
+            ]
+
+        self._last_logits = logits
+        return features, logits
+
+    def predict(
+        self,
+        dataset: Union[DataLoader, ItemType],
+        postproc_fns: Optional[List[Callable]] = None,
+        detach: bool = True,
+        verbose: bool = False,
+        numpy_concat: bool = False,
+        **kwargs,
+    ) -> Tuple[List[torch.Tensor], dict]:
+        """Get the projection of the dataset in the feature space of self.model
+
+        Args:
+            dataset (Union[DataLoader, ItemType]): input dataset
+            postproc_fns (Optional[List[Callable]]): postprocessing function to apply to
+                each feature immediately after forward. Default to None.
+            detach (bool): if True, return features detached from the computational
+                graph. No gradient will be computed. Defaults to True.
+            verbose (bool): if True, display a progress bar. Defaults to False.
+            numpy_concat (bool): if True, each mini-batch is immediately moved
+                to CPU and converted to a NumPy array before concatenation.
+                That keeps GPU memory constant at one batch, at the cost of a small
+                host-device transfer overhead. Defaults to False.
+
+        Returns:
+            List[torch.Tensor], dict: features and extra information (logits, labels) as
+                a dictionary.
+        """
+        labels = None
+
+        if isinstance(dataset, get_args(ItemType)):
+            tensor = TorchDataHandler.get_input_from_dataset_item(dataset)
+            features, logits = self.predict_tensor(tensor, postproc_fns, detach=detach)
+
+            # Get labels if dataset is a tuple/list
+            if isinstance(dataset, (list, tuple)) and len(dataset) > 1:
+                labels = TorchDataHandler.get_label_from_dataset_item(dataset)
+
+        else:
+            # Check if batches include labels
+            first_batch = next(iter(dataset))
+            contains_labels = (
+                isinstance(first_batch, (list, tuple)) and len(first_batch) > 1
+            )
+
+            # Prepare buffers
+            features_per_layer = [[] for _ in self.feature_layers_id]
+            logits_list = []
+            labels_list = [] if contains_labels else None
+
+            # Process batches
+            for elem in tqdm(dataset, desc="Predicting", disable=not verbose):
+                tensor = TorchDataHandler.get_input_from_dataset_item(elem)
+                feats_batch, logits_batch = self.predict_tensor(
+                    tensor, postproc_fns, detach=detach
+                )
+
+                # Move to host and convert if requested
+                if numpy_concat:
+                    feats_batch = [f.detach().cpu().numpy() for f in feats_batch]
+                    logits_batch = (
+                        logits_batch.detach().cpu().numpy()
+                        if logits_batch is not None
+                        else None
+                    )
+
+                # Accumulate
+                for i, f in enumerate(feats_batch):
+                    features_per_layer[i].append(f)
+                if logits_batch is not None:
+                    logits_list.append(logits_batch)
+                if contains_labels:
+                    lbl = TorchDataHandler.get_label_from_dataset_item(elem)
+                    labels_list.append(lbl)
+
+            # Concatenate
+            labels = torch.cat(labels_list, dim=0) if labels_list is not None else None
+
+            if numpy_concat:
+                features = [np.concatenate(lst, axis=0) for lst in features_per_layer]
+                logits = np.concatenate(logits_list, axis=0) if logits_list else None
+                labels = labels.cpu().numpy() if labels is not None else None
+
+            else:
+                features = [torch.cat(lst, dim=0) for lst in features_per_layer]
+                logits = torch.cat(logits_list, dim=0) if logits_list else None
+
+        # Package extra info
+        info = {"labels": labels, "logits": logits}
+        return features, info
+
+    def get_weights(self, layer_id: Union[str, int]) -> List[torch.Tensor]:
+        """Get the weights of a layer
+
+        Args:
+            layer_id (Union[int, str]): layer identifier
+
+        Returns:
+            List[torch.Tensor]: weights and biases matrixes
+        """
+        layer = self.find_layer(self.model, layer_id)
+        return [layer.weight.detach().cpu().numpy(), layer.bias.detach().cpu().numpy()]
+
+    def _get_clip_hook(self, threshold: float) -> Callable:
+        """
+        Hook that truncate activation features under a threshold value
+
+        Args:
+            threshold (float): threshold value
+
+        Returns:
+            Callable: hook function
+        """
+
+        def hook(_, input):
+            input = input[0]
+            input = torch.clip(input, max=threshold)
+            return input
+
+        return hook
+
+    def _get_scale_hook(self, percentile: float) -> Callable:
+        """
+        Hook that scales activation features.
+
+        Args:
+            threshold (float): threshold value
+
+        Returns:
+            Callable: hook function
+        """
+
+        def hook(_, input):
+            input = input[0]
+            output_percentile = torch.quantile(input, percentile, dim=1)
+            mask = input > output_percentile[:, None]
+            output_masked = input * mask
+            s = torch.exp(torch.sum(input, dim=1) / torch.sum(output_masked, dim=1))
+            s = torch.unsqueeze(s, 1)
+            input = input * s
+            return input
+
+        return hook
+
+    def _get_ash_hook(self, percentile: float) -> Callable:
+        """
+        Hook that scales and prunes activation features under a threshold value
+
+        Args:
+            threshold (float): threshold value
+
+        Returns:
+            Callable: hook function
+        """
+
+        def hook(_, input):
+            input = input[0]
+            output_percentile = torch.quantile(input, percentile, dim=1)
+            mask = input > output_percentile[:, None]
+            output_masked = input * mask
+            s = torch.exp(torch.sum(input, dim=1) / torch.sum(output_masked, dim=1))
+            s = torch.unsqueeze(s, 1)
+            input = output_masked * s
+            return 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 _default_postproc_fn(self, feat: TensorType) -> TensorType:
+        """Default postprocessing function to apply to each feature immediately
+        after forward.
+
+        Args:
+            feat (TensorType): input tensor
+
+        Returns:
+            TensorType: postprocessed tensor
+        """
+        if len(feat.shape) == 4:
+            feat = nn.AdaptiveAvgPool2d(1)(feat)
+            feat = feat.view(feat.size(0), -1)
+        elif len(feat.shape) == 3:
+            feat = nn.AdaptiveAvgPool1d(1)(feat)
+            feat = feat.view(feat.size(0), -1)
+        elif len(feat.shape) == 2:
+            feat = feat
+        else:
+            raise NotImplementedError(
+                "Postprocessing function not implemented for this feature shape"
+            )
+        return feat

oodeel/methods/__init__.py

@@ -0,0 +1,47 @@
+# -*- coding: utf-8 -*-
+# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+from .dknn import DKNN
+from .energy import Energy
+from .entropy import Entropy
+from .gen import GEN
+from .gram import Gram
+from .mahalanobis import Mahalanobis
+from .mls import MLS
+from .odin import ODIN
+from .rmds import RMDS
+from .she import SHE
+from .vim import VIM
+
+__all__ = [
+    "DKNN",
+    "Energy",
+    "Entropy",
+    "GEN",
+    "SHE",
+    "Gram",
+    "Mahalanobis",
+    "MLS",
+    "ODIN",
+    "RMDS",
+    "VIM",
+]