kaiko-eva 0.0.0.dev6__py3-none-any.whl

eva/core/models/modules/module.py

@@ -0,0 +1,190 @@
+"""Base model module."""
+
+from typing import Any, Mapping
+
+import lightning.pytorch as pl
+import torch
+from lightning.pytorch.utilities import memory
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+from typing_extensions import override
+
+from eva.core.metrics import structs as metrics_lib
+from eva.core.models.modules.typings import INPUT_BATCH
+from eva.core.models.modules.utils import batch_postprocess
+
+
+class ModelModule(pl.LightningModule):
+    """The base model module."""
+
+    def __init__(
+        self,
+        metrics: metrics_lib.MetricsSchema | None = None,
+        postprocess: batch_postprocess.BatchPostProcess | None = None,
+    ) -> None:
+        """Initializes the basic module.
+
+        Args:
+            metrics: The metric groups to track.
+            postprocess: A list of helper functions to apply after the
+                loss and before the metrics calculation to the model
+                predictions and targets.
+        """
+        super().__init__()
+
+        self._metrics = metrics or self.default_metrics
+        self._postprocess = postprocess or self.default_postprocess
+
+        self.metrics = metrics_lib.MetricModule.from_schema(self._metrics)
+
+    @property
+    def default_metrics(self) -> metrics_lib.MetricsSchema:
+        """The default metrics."""
+        return metrics_lib.MetricsSchema()
+
+    @property
+    def default_postprocess(self) -> batch_postprocess.BatchPostProcess:
+        """The default post-processes."""
+        return batch_postprocess.BatchPostProcess()
+
+    @override
+    def on_train_batch_end(
+        self,
+        outputs: STEP_OUTPUT,
+        batch: INPUT_BATCH,
+        batch_idx: int,
+    ) -> None:
+        outputs = self._common_batch_end(outputs)
+        self._forward_and_log_metrics(
+            self.metrics.training_metrics,
+            batch_outputs=outputs,
+        )
+
+    @override
+    def on_validation_batch_end(
+        self,
+        outputs: STEP_OUTPUT,
+        batch: INPUT_BATCH,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        outputs = self._common_batch_end(outputs)
+        self._update_metrics(
+            self.metrics.validation_metrics,
+            outputs=outputs,
+            dataloader_idx=dataloader_idx,
+        )
+
+    @override
+    def on_validation_epoch_end(self) -> None:
+        self._compute_and_log_metrics(self.metrics.validation_metrics)
+
+    @override
+    def on_test_batch_end(
+        self,
+        outputs: STEP_OUTPUT,
+        batch: INPUT_BATCH,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        outputs = self._common_batch_end(outputs)
+        self._update_metrics(
+            self.metrics.test_metrics,
+            outputs=outputs,
+            dataloader_idx=dataloader_idx,
+        )
+
+    @override
+    def on_test_epoch_end(self) -> None:
+        self._compute_and_log_metrics(self.metrics.test_metrics)
+
+    def _common_batch_end(self, outputs: STEP_OUTPUT) -> STEP_OUTPUT:
+        """Common end step of training, validation and test.
+
+        It will apply the post-processes to the batch output and move
+        them to the appropriate device.
+
+        Args:
+            outputs: The batch step outputs.
+
+        Returns:
+            The updated outputs.
+        """
+        self._postprocess(outputs)
+        return memory.recursive_detach(outputs, to_cpu=self.device.type == "cpu")
+
+    def _forward_and_log_metrics(
+        self,
+        metrics: metrics_lib.MetricCollection,
+        batch_outputs: STEP_OUTPUT,
+    ) -> None:
+        """Performs a forward pass to the metrics and logs them.
+
+        Args:
+            metrics: The desired metrics tracker to log.
+            batch_outputs: The outputs of the batch processing step.
+        """
+        inputs = self._parse_metrics_inputs(batch_outputs)
+        metrics(**inputs)
+        self.log_dict(metrics, on_step=True, on_epoch=False)
+
+    def _update_metrics(
+        self,
+        metrics: metrics_lib.MetricCollection,
+        outputs: STEP_OUTPUT,
+        dataloader_idx: int = 0,
+    ) -> None:
+        """Updates the metrics tracker with new data.
+
+        Here the `outputs` keyword values will be filtered based
+        on the signature of all individual metrics and passed only
+        to the compatible ones.
+
+        Args:
+            metrics: The desired metrics tracker to update.
+            outputs: The outputs of the batch processing step.
+            dataloader_idx: The dataloader index.
+        """
+        inputs = self._parse_metrics_inputs(outputs, dataloader_idx)
+        metrics.update(**inputs)
+
+    def _compute_and_log_metrics(self, metrics: metrics_lib.MetricCollection) -> None:
+        """Computes, logs and resets the metrics.
+
+        Args:
+            metrics: The desired metrics tracker to log.
+        """
+        outputs = metrics.compute()
+        self.log_dict(outputs)
+        metrics.reset()
+
+    def _parse_metrics_inputs(
+        self,
+        outputs: STEP_OUTPUT,
+        dataloader_idx: int = 0,
+    ) -> Mapping[str, Any]:
+        """Parses the arguments for the metrics.
+
+        When pass to a metrics collection object, the keyword values
+        will be filtered based on the signature of all individual
+        metrics and passed only to the compatible ones.
+
+        Args:
+            outputs: The outputs of the batch processing step.
+            dataloader_idx: The dataloader index.
+
+        Returns:
+            A mapping with the argument name and its value.
+        """
+        if outputs is None:
+            return {}
+
+        if isinstance(outputs, torch.Tensor):
+            outputs = {"loss": outputs}
+
+        additional_metric_inputs = {
+            "preds": outputs.get("predictions"),
+            "target": outputs.get("targets"),
+            "metadata": outputs.get("metadata"),
+            "dataloader_idx": dataloader_idx,
+        }
+        return {**additional_metric_inputs, **outputs}

eva/core/models/modules/typings.py

@@ -0,0 +1,23 @@
+"""Type annotations for model modules."""
+
+from typing import Any, Dict, NamedTuple
+
+import lightning.pytorch as pl
+import torch
+from torch import nn
+
+MODEL_TYPE = nn.Module | pl.LightningModule
+"""The expected model type."""
+
+
+class INPUT_BATCH(NamedTuple):
+    """The default input batch data scheme."""
+
+    data: torch.Tensor
+    """The data batch."""
+
+    targets: torch.Tensor | Dict[str, Any] | None = None
+    """The target batch."""
+
+    metadata: Dict[str, Any] | None = None
+    """The associated metadata."""

eva/core/models/modules/utils/__init__.py

@@ -0,0 +1,6 @@
+"""Utilities and helper functionalities for model modules."""
+
+from eva.core.models.modules.utils import grad
+from eva.core.models.modules.utils.batch_postprocess import BatchPostProcess
+
+__all__ = ["grad", "BatchPostProcess"]

eva/core/models/modules/utils/batch_postprocess.py

@@ -0,0 +1,57 @@
+"""Batch post-processes module."""
+
+import dataclasses
+import functools
+from typing import Callable, List
+
+import torch
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+
+Transform = Callable[[torch.Tensor], torch.Tensor]
+"""Post-process transform type."""
+
+
+@dataclasses.dataclass(frozen=True)
+class BatchPostProcess:
+    """Batch post-processes transform schema."""
+
+    targets_transforms: List[Transform] | None = None
+    """Holds the common train and evaluation metrics."""
+
+    predictions_transforms: List[Transform] | None = None
+    """Holds the common train and evaluation metrics."""
+
+    def __call__(self, outputs: STEP_OUTPUT) -> None:
+        """Applies the defined list of transforms to the batch output in-place.
+
+        Note that the transforms are applied only when the input is a dictionary
+        and only to its keys of `predictions` and/or `targets`.
+
+        Args:
+            outputs: The batch output of the model module step.
+        """
+        if not isinstance(outputs, dict):
+            return
+
+        if "targets" in outputs and self.targets_transforms is not None:
+            outputs["targets"] = _apply_transforms(
+                outputs["targets"], transforms=self.targets_transforms
+            )
+
+        if "predictions" in outputs and self.predictions_transforms is not None:
+            outputs["predictions"] = _apply_transforms(
+                outputs["predictions"], transforms=self.predictions_transforms
+            )
+
+
+def _apply_transforms(tensor: torch.Tensor, transforms: List[Transform]) -> torch.Tensor:
+    """Applies a list of transforms sequentially to a input tensor.
+
+    Args:
+        tensor: The desired tensor to process.
+        transforms: The list of transforms to apply to the input tensor.
+
+    Returns:
+        The processed tensor.
+    """
+    return functools.reduce(lambda tensor, transform: transform(tensor), transforms, tensor)

eva/core/models/modules/utils/grad.py

@@ -0,0 +1,23 @@
+"""Gradient related utilities and helper functions."""
+
+from torch import nn
+
+
+def deactivate_requires_grad(module: nn.Module) -> None:
+    """Deactivates the `requires_grad` flag for all parameters of a model.
+
+    Args:
+        module: The torch module to deactivate the gradient computation in place.
+    """
+    for parameter in module.parameters():
+        parameter.requires_grad = False
+
+
+def activate_requires_grad(module: nn.Module) -> None:
+    """Activates the `requires_grad` flag for all parameters of a model.
+
+    Args:
+        module: The torch module to deactivate the gradient computation in place.
+    """
+    for parameter in module.parameters():
+        parameter.requires_grad = True

eva/core/models/networks/__init__.py

@@ -0,0 +1,6 @@
+"""Networks API."""
+
+from eva.core.models.networks.mlp import MLP
+from eva.core.models.networks.wrappers import HuggingFaceModel, ModelFromFunction, ONNXModel
+
+__all__ = ["ModelFromFunction", "HuggingFaceModel", "ONNXModel", "MLP"]

eva/core/models/networks/_utils.py

@@ -0,0 +1,25 @@
+"""Utilities and helper functions for models."""
+
+from lightning_fabric.utilities import cloud_io
+from loguru import logger
+from torch import nn
+
+
+def load_model_weights(model: nn.Module, checkpoint_path: str) -> None:
+    """Loads (local or remote) weights to the model in-place.
+
+    Args:
+        model: The model to load the weights to.
+        checkpoint_path: The path to the model weights/checkpoint.
+    """
+    logger.info(f"Loading '{model.__class__.__name__}' model from checkpoint '{checkpoint_path}'")
+
+    fs = cloud_io.get_filesystem(checkpoint_path)
+    with fs.open(checkpoint_path, "rb") as file:
+        checkpoint = cloud_io._load(file, map_location="cpu")  # type: ignore
+        if isinstance(checkpoint, dict) and "state_dict" in checkpoint:
+            checkpoint = checkpoint["state_dict"]
+
+        model.load_state_dict(checkpoint, strict=True)
+
+    logger.info(f"Loading weights from '{checkpoint_path}' completed successfully.")

eva/core/models/networks/mlp.py

@@ -0,0 +1,69 @@
+"""Multi-layer Perceptron (MLP) implemented in PyTorch."""
+
+from typing import Type
+
+import torch
+import torch.nn as nn
+
+
+class MLP(nn.Module):
+    """A Multi-layer Perceptron (MLP) network."""
+
+    def __init__(
+        self,
+        input_size: int,
+        output_size: int,
+        hidden_layer_sizes: tuple[int, ...] | None = None,
+        hidden_activation_fn: Type[torch.nn.Module] | None = nn.ReLU,
+        output_activation_fn: Type[torch.nn.Module] | None = None,
+        dropout: float = 0.0,
+    ) -> None:
+        """Initializes the MLP.
+
+        Args:
+            input_size: The number of input features.
+            output_size: The number of output features.
+            hidden_layer_sizes: A list specifying the number of units in each hidden layer.
+            dropout: Dropout probability for hidden layers.
+            hidden_activation_fn: Activation function to use for hidden layers. Default is ReLU.
+            output_activation_fn: Activation function to use for the output layer. Default is None.
+        """
+        super().__init__()
+
+        self.input_size = input_size
+        self.output_size = output_size
+        self.hidden_layer_sizes = hidden_layer_sizes if hidden_layer_sizes is not None else ()
+        self.hidden_activation_fn = hidden_activation_fn
+        self.output_activation_fn = output_activation_fn
+        self.dropout = dropout
+
+        self._network = self._build_network()
+
+    def _build_network(self) -> nn.Sequential:
+        """Builds the neural network's layers and returns a nn.Sequential container."""
+        layers = []
+        prev_size = self.input_size
+        for size in self.hidden_layer_sizes:
+            layers.append(nn.Linear(prev_size, size))
+            if self.hidden_activation_fn is not None:
+                layers.append(self.hidden_activation_fn())
+            if self.dropout > 0:
+                layers.append(nn.Dropout(self.dropout))
+            prev_size = size
+
+        layers.append(nn.Linear(prev_size, self.output_size))
+        if self.output_activation_fn is not None:
+            layers.append(self.output_activation_fn())
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Defines the forward pass of the MLP.
+
+        Args:
+            x: The input tensor.
+
+        Returns:
+            The output of the network.
+        """
+        return self._network(x)

eva/core/models/networks/transforms/__init__.py

@@ -0,0 +1,5 @@
+"""Model outputs transforms API."""
+
+from eva.core.models.networks.transforms.extract_cls_features import ExtractCLSFeatures
+
+__all__ = ["ExtractCLSFeatures"]

eva/core/models/networks/transforms/extract_cls_features.py

@@ -0,0 +1,25 @@
+"""Transforms for extracting the CLS output from a model output."""
+
+import torch
+from transformers import modeling_outputs
+
+
+class ExtractCLSFeatures:
+    """Extracts the CLS token from a ViT model output."""
+
+    def __call__(
+        self, tensor: torch.Tensor | modeling_outputs.BaseModelOutputWithPooling
+    ) -> torch.Tensor:
+        """Call method for the transformation.
+
+        Args:
+            tensor: The tensor representing the model output.
+        """
+        if isinstance(tensor, torch.Tensor):
+            transformed_tensor = tensor[:, 0, :]
+        elif isinstance(tensor, modeling_outputs.BaseModelOutputWithPooling):
+            transformed_tensor = tensor.last_hidden_state[:, 0, :]
+        else:
+            raise ValueError(f"Unsupported type {type(tensor)}")
+
+        return transformed_tensor

eva/core/models/networks/wrappers/__init__.py

@@ -0,0 +1,8 @@
+"""Model Wrappers API."""
+
+from eva.core.models.networks.wrappers.base import BaseModel
+from eva.core.models.networks.wrappers.from_function import ModelFromFunction
+from eva.core.models.networks.wrappers.huggingface import HuggingFaceModel
+from eva.core.models.networks.wrappers.onnx import ONNXModel
+
+__all__ = ["BaseModel", "ModelFromFunction", "HuggingFaceModel", "ONNXModel"]

eva/core/models/networks/wrappers/base.py

@@ -0,0 +1,47 @@
+"""Base class for model wrappers."""
+
+import abc
+from typing import Callable
+
+import torch
+import torch.nn as nn
+from typing_extensions import override
+
+
+class BaseModel(nn.Module):
+    """Base class for model wrappers."""
+
+    def __init__(self, tensor_transforms: Callable | None = None) -> None:
+        """Initializes the model.
+
+        Args:
+            tensor_transforms: The transforms to apply to the output
+                tensor produced by the model.
+        """
+        super().__init__()
+
+        self._output_transforms = tensor_transforms
+
+    @override
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        tensor = self.model_forward(tensor)
+        return self._apply_transforms(tensor)
+
+    @abc.abstractmethod
+    def load_model(self) -> Callable[..., torch.Tensor]:
+        """Loads the model."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def model_forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Implements the forward pass of the model.
+
+        Args:
+            tensor: The input tensor to the model.
+        """
+        raise NotImplementedError
+
+    def _apply_transforms(self, tensor: torch.Tensor) -> torch.Tensor:
+        if self._output_transforms is not None:
+            tensor = self._output_transforms(tensor)
+        return tensor

eva/core/models/networks/wrappers/from_function.py

@@ -0,0 +1,58 @@
+"""Helper function from models defined with a function."""
+
+from typing import Any, Callable, Dict
+
+import jsonargparse
+import torch
+from torch import nn
+from typing_extensions import override
+
+from eva.core.models.networks import _utils
+from eva.core.models.networks.wrappers import base
+
+
+class ModelFromFunction(base.BaseModel):
+    """Wrapper class for models which are initialized from functions.
+
+    This is helpful for initializing models in a `.yaml` configuration file.
+    """
+
+    def __init__(
+        self,
+        path: Callable[..., nn.Module],
+        arguments: Dict[str, Any] | None = None,
+        checkpoint_path: str | None = None,
+        tensor_transforms: Callable | None = None,
+    ) -> None:
+        """Initializes and constructs the model.
+
+        Args:
+            path: The path to the callable object (class or function).
+            arguments: The extra callable function / class arguments.
+            checkpoint_path: The path to the checkpoint to load the model
+                weights from. This is currently only supported for torch
+                model checkpoints. For other formats, the checkpoint loading
+                should be handled within the provided callable object in <path>.
+            tensor_transforms: The transforms to apply to the output tensor
+                produced by the model.
+        """
+        super().__init__()
+
+        self._path = path
+        self._arguments = arguments
+        self._checkpoint_path = checkpoint_path
+        self._tensor_transforms = tensor_transforms
+
+        self._model = self.load_model()
+
+    @override
+    def load_model(self) -> nn.Module:
+        class_path = jsonargparse.class_from_function(self._path, func_return=nn.Module)
+        model = class_path(**self._arguments or {})
+        if self._checkpoint_path is not None:
+            _utils.load_model_weights(model, self._checkpoint_path)
+        return model
+
+    @override
+    def model_forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        return self._model(tensor)

eva/core/models/networks/wrappers/huggingface.py

@@ -0,0 +1,37 @@
+"""Wrappers for HuggingFace `transformers` models."""
+
+from typing import Any, Callable
+
+import torch
+import transformers
+from typing_extensions import override
+
+from eva.core.models.networks.wrappers import base
+
+
+class HuggingFaceModel(base.BaseModel):
+    """Wrapper class for loading HuggingFace `transformers` models."""
+
+    def __init__(self, model_name_or_path: str, tensor_transforms: Callable | None = None) -> None:
+        """Initializes the model.
+
+        Args:
+            model_name_or_path: The model name or path to load the model from.
+                This can be a local path or a model name from the `HuggingFace`
+                model hub.
+            tensor_transforms: The transforms to apply to the output tensor
+                produced by the model.
+        """
+        super().__init__(tensor_transforms=tensor_transforms)
+
+        self._model_name_or_path = model_name_or_path
+        self._model = self.load_model()
+
+    @override
+    def load_model(self) -> Any:
+        config = transformers.AutoConfig.from_pretrained(self._model_name_or_path)
+        return transformers.AutoModel.from_pretrained(self._model_name_or_path, config=config)
+
+    @override
+    def model_forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        return self._model(tensor)

eva/core/models/networks/wrappers/onnx.py

@@ -0,0 +1,47 @@
+"""Wrapper class for ONNX models."""
+
+from typing import Any, Callable, Literal
+
+import onnxruntime as ort
+import torch
+from typing_extensions import override
+
+from eva.core.models.networks.wrappers import base
+
+
+class ONNXModel(base.BaseModel):
+    """Wrapper class for loading ONNX models."""
+
+    def __init__(
+        self,
+        path: str,
+        device: Literal["cpu", "cuda"] | None = "cpu",
+        tensor_transforms: Callable | None = None,
+    ):
+        """Initializes the model.
+
+        Args:
+            path: The path to the .onnx model file.
+            device: The device to run the model on. This can be either "cpu" or "cuda".
+            tensor_transforms: The transforms to apply to the output tensor produced by the model.
+        """
+        super().__init__(tensor_transforms=tensor_transforms)
+
+        self._path = path
+        self._device = device
+        self._model = self.load_model()
+
+    @override
+    def load_model(self) -> Any:
+        if self._device == "cuda" and not torch.cuda.is_available():
+            raise ValueError("Device is set to 'cuda', but CUDA is not available.")
+        provider = "CUDAExecutionProvider" if self._device == "cuda" else "CPUExecutionProvider"
+        return ort.InferenceSession(self._path, providers=[provider])
+
+    @override
+    def model_forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        # TODO: Use IO binding to avoid copying the tensor to CPU.
+        # https://onnxruntime.ai/docs/api/python/api_summary.html#data-on-device
+        inputs = {self._model.get_inputs()[0].name: tensor.detach().cpu().numpy()}
+        outputs = self._model.run(None, inputs)[0]
+        return torch.from_numpy(outputs).float().to(tensor.device)

eva/core/trainers/__init__.py

@@ -0,0 +1,6 @@
+"""Trainers API."""
+
+from eva.core.trainers.functional import infer_model, run_evaluation_session
+from eva.core.trainers.trainer import Trainer
+
+__all__ = ["infer_model", "run_evaluation_session", "Trainer"]