quadra 0.0.1__py3-none-any.whl → 2.1.13__py3-none-any.whl

quadra/models/base.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+import inspect
+from collections.abc import Sequence
+from typing import Any
+
+import torch
+from torch import nn
+
+from quadra.utils.logger import get_logger
+
+log = get_logger(__name__)
+
+
+class ModelSignatureWrapper(nn.Module):
+    """Model wrapper used to retrieve input shape. It can be used as a decorator of nn.Module, the first call to the
+    forward method will retrieve the input shape and store it in the input_shapes attribute.
+    It will also save the model summary in a file called model_summary.txt in the current working directory.
+    """
+
+    def __init__(self, model: nn.Module):
+        super().__init__()
+        self.instance = model
+        self.input_shapes: Any = None
+        self.disable = False
+
+        if isinstance(self.instance, ModelSignatureWrapper):
+            # Handle nested ModelSignatureWrapper
+            self.input_shapes = self.instance.input_shapes
+            self.instance = self.instance.instance
+
+    def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
+        """Retrieve the input shape and forward the model, if the input shape is already retrieved it will just forward
+        the model.
+        """
+        if self.input_shapes is None and not self.disable:
+            try:
+                self.input_shapes = self._get_input_shapes(*args, **kwargs)
+            except Exception:
+                log.warning(
+                    "Failed to retrieve input shapes after forward! To export the model you'll need to "
+                    "provide the input shapes manually setting the config.export.input_shapes parameter! "
+                    "Alternatively you could try to use a forward with supported input types (and their compositions) "
+                    "(list, tuple, dict, tensors)."
+                )
+                self.disable = True
+
+        return self.instance.forward(*args, **kwargs)
+
+    def to(self, *args, **kwargs):
+        """Handle calls to to method returning the underlying model."""
+        self.instance = self.instance.to(*args, **kwargs)
+
+        return self
+
+    def half(self, *args, **kwargs):
+        """Handle calls to to method returning the underlying model."""
+        self.instance = self.instance.half(*args, **kwargs)
+
+        return self
+
+    def cpu(self, *args, **kwargs):
+        """Handle calls to to method returning the underlying model."""
+        self.instance = self.instance.cpu(*args, **kwargs)
+
+        return self
+
+    def _get_input_shapes(self, *args: Any, **kwargs: Any) -> list[Any]:
+        """Retrieve the input shapes from the input. Inputs will be in the same order as the forward method
+        signature.
+        """
+        input_shapes = []
+
+        for arg in args:
+            input_shapes.append(self._get_input_shape(arg))
+
+        if isinstance(self.instance.forward, torch.ScriptMethod):
+            # Handle torchscript backbones
+            for i, argument in enumerate(self.instance.forward.schema.arguments):
+                if i < (len(args) + 1):  # +1 for self
+                    continue
+
+                if argument.name == "self":
+                    continue
+
+                if argument.name in kwargs:
+                    input_shapes.append(self._get_input_shape(kwargs[argument.name]))
+                else:
+                    # Retrieve the default value
+                    input_shapes.append(self._get_input_shape(argument.default_value))
+        else:
+            signature = inspect.signature(self.instance.forward)
+
+            for i, key in enumerate(signature.parameters.keys()):
+                if i < len(args):
+                    continue
+
+                if key in kwargs:
+                    input_shapes.append(self._get_input_shape(kwargs[key]))
+                else:
+                    # Retrieve the default value
+                    input_shapes.append(self._get_input_shape(signature.parameters[key].default))
+
+        return input_shapes
+
+    def _get_input_shape(self, inp: Sequence | torch.Tensor) -> list[Any] | tuple[Any, ...] | dict[str, Any]:
+        """Recursive function to retrieve the input shapes."""
+        if isinstance(inp, list):
+            return [self._get_input_shape(i) for i in inp]
+
+        if isinstance(inp, tuple):
+            return tuple(self._get_input_shape(i) for i in inp)
+
+        if isinstance(inp, torch.Tensor):
+            return tuple(inp.shape[1:])
+
+        if isinstance(inp, dict):
+            return {k: self._get_input_shape(v) for k, v in inp.items()}
+
+        raise ValueError(f"Input type {type(inp)} not supported")
+
+    def __getattr__(self, name: str) -> torch.Tensor | nn.Module:
+        if name in ["instance", "input_shapes"]:
+            return self.__dict__[name]
+
+        return getattr(self.__dict__["instance"], name)
+
+    def __setattr__(self, name: str, value: torch.Tensor | nn.Module) -> None:
+        if name in ["instance", "input_shapes"]:
+            self.__dict__[name] = value
+        else:
+            setattr(self.instance, name, value)
+
+    def __getattribute__(self, __name: str) -> Any:
+        if __name in [
+            "instance",
+            "input_shapes",
+            "__dict__",
+            "forward",
+            "_get_input_shapes",
+            "_get_input_shape",
+            "to",
+            "half",
+            "cpu",
+            "call_super_init",
+            "_call_impl",
+            "_compiled_call_impl",
+        ]:
+            return super().__getattribute__(__name)
+
+        return getattr(self.instance, __name)

quadra/models/classification/__init__.py

@@ -0,0 +1,8 @@
+from .backbones import BaseNetworkBuilder, TimmNetworkBuilder, TorchHubNetworkBuilder, TorchVisionNetworkBuilder
+
+__all__ = [
+    "BaseNetworkBuilder",
+    "TorchVisionNetworkBuilder",
+    "TorchHubNetworkBuilder",
+    "TimmNetworkBuilder",
+]

quadra/models/classification/backbones.py

@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+from typing import Any
+
+import timm
+import torch
+from timm.models.helpers import load_checkpoint
+from torch import nn
+from torchvision import models
+
+from quadra.models.classification.base import BaseNetworkBuilder
+from quadra.utils.logger import get_logger
+
+log = get_logger(__name__)
+
+
+class TorchHubNetworkBuilder(BaseNetworkBuilder):
+    """TorchHub feature extractor, with the possibility to map features to an hypersphere.
+
+    Args:
+        repo_or_dir: The name of the repository or the path to the directory containing the model.
+        model_name: The name of the model within the repository.
+        pretrained: Whether to load the pretrained weights for the model.
+        pre_classifier: Pre classifier as a torch.nn.Module. Defaults to nn.Identity() if None.
+        classifier: Classifier as a torch.nn.Module. Defaults to nn.Identity() if None.
+        freeze: Whether to freeze the feature extractor. Defaults to True.
+        hyperspherical: Whether to map features to an hypersphere. Defaults to False.
+        flatten_features: Whether to flatten the features before the pre_classifier. Defaults to True.
+        checkpoint_path: Path to a checkpoint to load after the model is initialized. Defaults to None.
+        **torch_hub_kwargs: Additional arguments to pass to torch.hub.load
+    """
+
+    def __init__(
+        self,
+        repo_or_dir: str,
+        model_name: str,
+        pretrained: bool = True,
+        pre_classifier: nn.Module | None = None,
+        classifier: nn.Module | None = None,
+        freeze: bool = True,
+        hyperspherical: bool = False,
+        flatten_features: bool = True,
+        checkpoint_path: str | None = None,
+        **torch_hub_kwargs: Any,
+    ):
+        self.pretrained = pretrained
+        features_extractor = torch.hub.load(
+            repo_or_dir=repo_or_dir, model=model_name, pretrained=self.pretrained, **torch_hub_kwargs
+        )
+        if checkpoint_path:
+            log.info("Loading checkpoint from %s", checkpoint_path)
+            load_checkpoint(model=features_extractor, checkpoint_path=checkpoint_path)
+
+        super().__init__(
+            features_extractor=features_extractor,
+            pre_classifier=pre_classifier,
+            classifier=classifier,
+            freeze=freeze,
+            hyperspherical=hyperspherical,
+            flatten_features=flatten_features,
+        )
+
+
+class TorchVisionNetworkBuilder(BaseNetworkBuilder):
+    """Torchvision feature extractor, with the possibility to map features to an hypersphere.
+
+    Args:
+        model_name: Torchvision model function that will be evaluated, for example: torchvision.models.resnet18.
+        pretrained: Whether to load the pretrained weights for the model.
+        pre_classifier: Pre classifier as a torch.nn.Module. Defaults to nn.Identity() if None.
+        classifier: Classifier as a torch.nn.Module. Defaults to nn.Identity() if None.
+        freeze: Whether to freeze the feature extractor. Defaults to True.
+        hyperspherical: Whether to map features to an hypersphere. Defaults to False.
+        flatten_features: Whether to flatten the features before the pre_classifier. Defaults to True.
+        checkpoint_path: Path to a checkpoint to load after the model is initialized. Defaults to None.
+        **torchvision_kwargs: Additional arguments to pass to the model function.
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        pretrained: bool = True,
+        pre_classifier: nn.Module | None = None,
+        classifier: nn.Module | None = None,
+        freeze: bool = True,
+        hyperspherical: bool = False,
+        flatten_features: bool = True,
+        checkpoint_path: str | None = None,
+        **torchvision_kwargs: Any,
+    ):
+        self.pretrained = pretrained
+        model_function = models.__dict__[model_name]
+        features_extractor = model_function(pretrained=self.pretrained, progress=True, **torchvision_kwargs)
+        if checkpoint_path:
+            log.info("Loading checkpoint from %s", checkpoint_path)
+            load_checkpoint(model=features_extractor, checkpoint_path=checkpoint_path)
+
+        # Remove classifier
+        features_extractor.classifier = nn.Identity()
+        super().__init__(
+            features_extractor=features_extractor,
+            pre_classifier=pre_classifier,
+            classifier=classifier,
+            freeze=freeze,
+            hyperspherical=hyperspherical,
+            flatten_features=flatten_features,
+        )
+
+
+class TimmNetworkBuilder(BaseNetworkBuilder):
+    """Torchvision feature extractor, with the possibility to map features to an hypersphere.
+
+    Args:
+        model_name: Timm model name
+        pretrained: Whether to load the pretrained weights for the model.
+        pre_classifier: Pre classifier as a torch.nn.Module. Defaults to nn.Identity() if None.
+        classifier: Classifier as a torch.nn.Module. Defaults to nn.Identity() if None.
+        freeze: Whether to freeze the feature extractor. Defaults to True.
+        hyperspherical: Whether to map features to an hypersphere. Defaults to False.
+        flatten_features: Whether to flatten the features before the pre_classifier. Defaults to True.
+        checkpoint_path: Path to a checkpoint to load after the model is initialized. Defaults to None.
+        **timm_kwargs: Additional arguments to pass to timm.create_model
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        pretrained: bool = True,
+        pre_classifier: nn.Module | None = None,
+        classifier: nn.Module | None = None,
+        freeze: bool = True,
+        hyperspherical: bool = False,
+        flatten_features: bool = True,
+        checkpoint_path: str | None = None,
+        **timm_kwargs: Any,
+    ):
+        self.pretrained = pretrained
+        features_extractor = timm.create_model(
+            model_name, pretrained=self.pretrained, num_classes=0, checkpoint_path=checkpoint_path, **timm_kwargs
+        )
+
+        super().__init__(
+            features_extractor=features_extractor,
+            pre_classifier=pre_classifier,
+            classifier=classifier,
+            freeze=freeze,
+            hyperspherical=hyperspherical,
+            flatten_features=flatten_features,
+        )

quadra/models/classification/base.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+from torch import nn
+
+from quadra.utils.models import L2Norm
+
+
+class BaseNetworkBuilder(nn.Module):
+    """Baseline Feature Extractor, with the possibility to map features to an hypersphere.
+        If hypershperical is True the classifier is ignored.
+
+    Args:
+        features_extractor: Feature extractor as a toch.nn.Module.
+        pre_classifier: Pre classifier as a torch.nn.Module. Defaults to nn.Identity() if None.
+        classifier: Classifier as a torch.nn.Module. Defaults to nn.Identity() if None.
+        freeze: Whether to freeze the feature extractor. Defaults to True.
+        hyperspherical: Whether to map features to an hypersphere. Defaults to False.
+        flatten_features: Whether to flatten the features before the pre_classifier. May be required if your model
+            is outputting a feature map rather than a vector. Defaults to True.
+    """
+
+    def __init__(
+        self,
+        features_extractor: nn.Module,
+        pre_classifier: nn.Module | None = None,
+        classifier: nn.Module | None = None,
+        freeze: bool = True,
+        hyperspherical: bool = False,
+        flatten_features: bool = True,
+    ):
+        super().__init__()
+        if pre_classifier is None:
+            pre_classifier = nn.Identity()
+
+        if classifier is None:
+            classifier = nn.Identity()
+
+        self.features_extractor = features_extractor
+        self.freeze = freeze
+        self.hyperspherical = hyperspherical
+        self.pre_classifier = pre_classifier
+        self.classifier = classifier
+        self.flatten: bool = False
+        self._hyperspherical: bool = False
+        self.l2: L2Norm | None = None
+        self.flatten_features = flatten_features
+
+        self.freeze = freeze
+        self.hyperspherical = hyperspherical
+
+        if self.freeze:
+            for p in self.features_extractor.parameters():
+                p.requires_grad = False
+
+    @property
+    def freeze(self) -> bool:
+        """Whether to freeze the feature extractor."""
+        return self._freeze
+
+    @freeze.setter
+    def freeze(self, value: bool) -> None:
+        """Whether to freeze the feature extractor."""
+        for p in self.features_extractor.parameters():
+            p.requires_grad = not value
+
+        self._freeze = value
+
+    @property
+    def hyperspherical(self) -> bool:
+        """Whether to map the extracted features into an hypersphere."""
+        return self._hyperspherical
+
+    @hyperspherical.setter
+    def hyperspherical(self, value: bool) -> None:
+        """Whether to map the extracted features into an hypersphere."""
+        self._hyperspherical = value
+        self.l2 = L2Norm() if value else None
+
+    def forward(self, x):
+        x = self.features_extractor(x)
+
+        if self.flatten_features:
+            x = x.view(x.size(0), -1)
+
+        x = self.pre_classifier(x)
+
+        if self.hyperspherical:
+            x = self.l2(x)
+
+        x = self.classifier(x)
+
+        return x