kaiko-eva 0.0.2__py3-none-any.whl → 0.1.0__py3-none-any.whl

eva/core/models/__init__.py

@@ -1,12 +1,14 @@
 """Models API."""
 
 from eva.core.models.modules import HeadModule, InferenceModule
-from eva.core.models.networks import MLP, HuggingFaceModel, ModelFromFunction, ONNXModel
+from eva.core.models.networks import MLP
+from eva.core.models.wrappers import BaseModel, HuggingFaceModel, ModelFromFunction, ONNXModel
 
 __all__ = [
     "HeadModule",
     "InferenceModule",
     "MLP",
+    "BaseModel",
     "HuggingFaceModel",
     "ModelFromFunction",
     "ONNXModel",

eva/core/models/modules/head.py

@@ -1,11 +1,11 @@
 """"Neural Network Head Module."""
 
-from typing import Any, Callable
+from typing import Any, Callable, Dict
 
 import torch
 from lightning.pytorch.cli import LRSchedulerCallable, OptimizerCallable
 from lightning.pytorch.utilities.types import STEP_OUTPUT
-from torch import optim
+from torch import nn, optim
 from torch.optim import lr_scheduler
 from typing_extensions import override
 
@@ -13,6 +13,7 @@ from eva.core.metrics import structs as metrics_lib
 from eva.core.models.modules import module
 from eva.core.models.modules.typings import INPUT_BATCH, MODEL_TYPE
 from eva.core.models.modules.utils import batch_postprocess, grad
+from eva.core.utils import parser
 
 
 class HeadModule(module.ModelModule):
@@ -24,7 +25,7 @@ class HeadModule(module.ModelModule):
 
     def __init__(
         self,
-        head: MODEL_TYPE,
+        head: Dict[str, Any] | MODEL_TYPE,
         criterion: Callable[..., torch.Tensor],
         backbone: MODEL_TYPE | None = None,
         optimizer: OptimizerCallable = optim.Adam,
@@ -36,6 +37,8 @@ class HeadModule(module.ModelModule):
 
         Args:
             head: The neural network that would be trained on the features.
+                If its a dictionary, it will be parsed to an object during the
+                `configure_model` step.
             criterion: The loss function to use.
             backbone: The feature extractor. If `None`, it will be expected
                 that the input batch returns the features directly.
@@ -48,7 +51,7 @@ class HeadModule(module.ModelModule):
         """
         super().__init__(metrics=metrics, postprocess=postprocess)
 
-        self.head = head
+        self.head = head  # type: ignore
         self.criterion = criterion
         self.backbone = backbone
         self.optimizer = optimizer
@@ -59,6 +62,9 @@ class HeadModule(module.ModelModule):
         if self.backbone is not None:
             grad.deactivate_requires_grad(self.backbone)
 
+        if isinstance(self.head, dict):
+            self.head: MODEL_TYPE = parser.parse_object(self.head, expected_type=nn.Module)
+
     @override
     def configure_optimizers(self) -> Any:
         parameters = self.head.parameters()

eva/core/models/modules/typings.py

@@ -16,7 +16,20 @@ class INPUT_BATCH(NamedTuple):
     data: torch.Tensor
     """The data batch."""
 
-    targets: torch.Tensor | Dict[str, Any] | None = None
+    targets: torch.Tensor | None = None
+    """The target batch."""
+
+    metadata: Dict[str, Any] | None = None
+    """The associated metadata."""
+
+
+class INPUT_TENSOR_BATCH(NamedTuple):
+    """Tensor based input batch data scheme."""
+
+    data: torch.Tensor
+    """The data batch."""
+
+    targets: torch.Tensor
     """The target batch."""
 
     metadata: Dict[str, Any] | None = None

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

@@ -2,9 +2,10 @@
 
 import dataclasses
 import functools
-from typing import Callable, List
+from typing import Any, Callable, Dict, List
 
 import torch
+from jsonargparse import _util
 from lightning.pytorch.utilities.types import STEP_OUTPUT
 
 Transform = Callable[[torch.Tensor], torch.Tensor]
@@ -15,10 +16,10 @@ Transform = Callable[[torch.Tensor], torch.Tensor]
 class BatchPostProcess:
     """Batch post-processes transform schema."""
 
-    targets_transforms: List[Transform] | None = None
+    targets_transforms: List[Transform | Dict[str, Any]] | None = None
     """Holds the common train and evaluation metrics."""
 
-    predictions_transforms: List[Transform] | None = None
+    predictions_transforms: List[Transform | Dict[str, Any]] | None = None
     """Holds the common train and evaluation metrics."""
 
     def __call__(self, outputs: STEP_OUTPUT) -> None:
@@ -35,12 +36,13 @@ class BatchPostProcess:
 
         if "targets" in outputs and self.targets_transforms is not None:
             outputs["targets"] = _apply_transforms(
-                outputs["targets"], transforms=self.targets_transforms
+                outputs["targets"], transforms=_parse_callable_inputs(self.targets_transforms)
             )
 
         if "predictions" in outputs and self.predictions_transforms is not None:
             outputs["predictions"] = _apply_transforms(
-                outputs["predictions"], transforms=self.predictions_transforms
+                outputs["predictions"],
+                transforms=_parse_callable_inputs(self.predictions_transforms),
             )
 
 
@@ -55,3 +57,33 @@ def _apply_transforms(tensor: torch.Tensor, transforms: List[Transform]) -> torc
         The processed tensor.
     """
     return functools.reduce(lambda tensor, transform: transform(tensor), transforms, tensor)
+
+
+def _parse_callable_inputs(inputs: List[Callable | Dict[str, Any]]) -> List[Callable]:
+    """Parses the inputs which where passed as dictionary to callable objects."""
+    parsed = []
+    for item in inputs:
+        if isinstance(item, dict):
+            item = _parse_dict(item)
+        parsed.append(item)
+    return parsed
+
+
+def _parse_dict(item: Dict[str, Any]) -> Callable:
+    """Parses the input dictionary to a partial callable object."""
+    if not _is_valid_dict(item):
+        raise ValueError(
+            "Transform dictionary format is not valid. "
+            "It must contain a key 'class_path' and optionally 'init_args' for "
+            "the function and additional call arguments."
+        )
+
+    return functools.partial(
+        _util.import_object(item["class_path"]),
+        **item.get("init_args", {}),
+    )
+
+
+def _is_valid_dict(item: Dict[str, Any], /) -> bool:
+    """Checks if the input has the valid structure."""
+    return "class_path" in item and set(item.keys()) <= {"class_path", "init_args"}

eva/core/models/networks/__init__.py

@@ -1,6 +1,5 @@
 """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"]
+__all__ = ["MLP"]

eva/core/models/networks/mlp.py

@@ -1,6 +1,6 @@
 """Multi-layer Perceptron (MLP) implemented in PyTorch."""
 
-from typing import Type
+from typing import Tuple, Type
 
 import torch
 import torch.nn as nn
@@ -13,7 +13,7 @@ class MLP(nn.Module):
         self,
         input_size: int,
         output_size: int,
-        hidden_layer_sizes: tuple[int, ...] | None = None,
+        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,

eva/core/models/transforms/__init__.py

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

eva/core/models/{networks/transforms → transforms}/extract_cls_features.py

@@ -7,6 +7,14 @@ from transformers import modeling_outputs
 class ExtractCLSFeatures:
     """Extracts the CLS token from a ViT model output."""
 
+    def __init__(self, cls_index: int = 0) -> None:
+        """Initializes the transformation.
+
+        Args:
+            cls_index: The index of the CLS token in the output tensor.
+        """
+        self._cls_index = cls_index
+
     def __call__(
         self, tensor: torch.Tensor | modeling_outputs.BaseModelOutputWithPooling
     ) -> torch.Tensor:
@@ -16,9 +24,9 @@ class ExtractCLSFeatures:
             tensor: The tensor representing the model output.
         """
         if isinstance(tensor, torch.Tensor):
-            transformed_tensor = tensor[:, 0, :]
+            transformed_tensor = tensor[:, self._cls_index, :]
         elif isinstance(tensor, modeling_outputs.BaseModelOutputWithPooling):
-            transformed_tensor = tensor.last_hidden_state[:, 0, :]
+            transformed_tensor = tensor.last_hidden_state[:, self._cls_index, :]
         else:
             raise ValueError(f"Unsupported type {type(tensor)}")
 

eva/core/models/transforms/extract_patch_features.py

@@ -0,0 +1,47 @@
+"""Transforms for extracting the patch features from a model output."""
+
+import math
+from typing import List
+
+import torch
+from transformers import modeling_outputs
+
+
+class ExtractPatchFeatures:
+    """Extracts the patch features from a ViT model output."""
+
+    def __init__(self, ignore_remaining_dims: bool = False) -> None:
+        """Initializes the transformation.
+
+        Args:
+            ignore_remaining_dims: If set to `True`, ignore the remaining dimensions
+                of the patch grid if it is not a square number.
+        """
+        self._ignore_remaining_dims = ignore_remaining_dims
+
+    def __call__(
+        self, tensor: torch.Tensor | modeling_outputs.BaseModelOutputWithPooling
+    ) -> List[torch.Tensor]:
+        """Call method for the transformation.
+
+        Args:
+            tensor: The raw embeddings of the model.
+
+        Returns:
+            A tensor (batch_size, hidden_size, n_patches_height, n_patches_width)
+            representing the model output.
+        """
+        if isinstance(tensor, modeling_outputs.BaseModelOutputWithPooling):
+            features = tensor.last_hidden_state[:, 1:, :].permute(0, 2, 1)
+            batch_size, hidden_size, patch_grid = features.shape
+            height = width = int(math.sqrt(patch_grid))
+            if height * width != patch_grid:
+                if self._ignore_remaining_dims:
+                    features = features[:, :, : height * width]
+                else:
+                    raise ValueError(f"Patch grid size must be a square number {patch_grid}.")
+            patch_embeddings = features.view(batch_size, hidden_size, height, width)
+        else:
+            raise ValueError(f"Unsupported type {type(tensor)}")
+
+        return [patch_embeddings]

eva/core/models/wrappers/__init__.py

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

eva/core/models/{networks/wrappers → wrappers}/base.py

@@ -22,6 +22,8 @@ class BaseModel(nn.Module):
 
         self._output_transforms = tensor_transforms
 
+        self._model: Callable[..., torch.Tensor] | nn.Module
+
     @override
     def forward(self, tensor: torch.Tensor) -> torch.Tensor:
         tensor = self.model_forward(tensor)
@@ -32,14 +34,13 @@ class BaseModel(nn.Module):
         """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
+        return self._model(tensor)
 
     def _apply_transforms(self, tensor: torch.Tensor) -> torch.Tensor:
         if self._output_transforms is not None:

eva/core/models/{networks/wrappers → wrappers}/from_function.py

@@ -3,12 +3,10 @@
 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
+from eva.core.models.wrappers import _utils, base
 
 
 class ModelFromFunction(base.BaseModel):
@@ -36,23 +34,18 @@ class ModelFromFunction(base.BaseModel):
             tensor_transforms: The transforms to apply to the output tensor
                 produced by the model.
         """
-        super().__init__()
+        super().__init__(tensor_transforms=tensor_transforms)
 
         self._path = path
         self._arguments = arguments
         self._checkpoint_path = checkpoint_path
-        self._tensor_transforms = tensor_transforms
 
-        self._model = self.load_model()
+        self.load_model()
 
     @override
-    def load_model(self) -> nn.Module:
+    def load_model(self) -> None:
         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)
+        self._model = model

eva/core/models/{networks/wrappers → wrappers}/huggingface.py

@@ -1,18 +1,22 @@
 """Wrappers for HuggingFace `transformers` models."""
 
-from typing import Any, Callable
+from typing import Any, Callable, Dict
 
-import torch
 import transformers
 from typing_extensions import override
 
-from eva.core.models.networks.wrappers import base
+from eva.core.models.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:
+    def __init__(
+        self,
+        model_name_or_path: str,
+        tensor_transforms: Callable | None = None,
+        model_kwargs: Dict[str, Any] | None = None,
+    ) -> None:
         """Initializes the model.
 
         Args:
@@ -21,17 +25,17 @@ class HuggingFaceModel(base.BaseModel):
                 model hub.
             tensor_transforms: The transforms to apply to the output tensor
                 produced by the model.
+            model_kwargs: The arguments used for instantiating the model.
         """
         super().__init__(tensor_transforms=tensor_transforms)
 
         self._model_name_or_path = model_name_or_path
-        self._model = self.load_model()
+        self._model_kwargs = model_kwargs or {}
 
-    @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)
+        self.load_model()
 
     @override
-    def model_forward(self, tensor: torch.Tensor) -> torch.Tensor:
-        return self._model(tensor)
+    def load_model(self) -> None:
+        self._model = transformers.AutoModel.from_pretrained(
+            self._model_name_or_path, **self._model_kwargs
+        )

eva/core/models/{networks/wrappers → wrappers}/onnx.py

@@ -6,7 +6,7 @@ import onnxruntime as ort
 import torch
 from typing_extensions import override
 
-from eva.core.models.networks.wrappers import base
+from eva.core.models.wrappers import base
 
 
 class ONNXModel(base.BaseModel):
@@ -29,19 +29,22 @@ class ONNXModel(base.BaseModel):
 
         self._path = path
         self._device = device
-        self._model = self.load_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])
+        self._model = ort.InferenceSession(self._path, providers=[provider])  # type: ignore
 
     @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
+        if not isinstance(self._model, ort.InferenceSession):
+            raise ValueError("Model is not loaded.")
         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/functional.py

@@ -69,6 +69,7 @@ def run_evaluation(
         A tuple of with the validation and the test metrics (if exists).
     """
     trainer, model = _utils.clone(base_trainer, base_model)
+    model.configure_model()
     trainer.setup_log_dirs(run_id or "")
     return fit_and_validate(trainer, model, datamodule, verbose=verbose)
 

eva/core/utils/__init__.py

@@ -1 +1,7 @@
 """Utilities and library level helper functionalities."""
+
+from eva.core.utils.clone import clone
+from eva.core.utils.memory import to_cpu
+from eva.core.utils.operations import numeric_sort
+
+__all__ = ["clone", "to_cpu", "numeric_sort"]

eva/core/utils/clone.py

@@ -0,0 +1,27 @@
+"""Clone related functions."""
+
+import functools
+from typing import Any, Dict, List
+
+import torch
+
+
+@functools.singledispatch
+def clone(tensor_type: Any) -> Any:
+    """Clone tensor objects."""
+    raise TypeError(f"Unsupported input type: {type(input)}.")
+
+
+@clone.register
+def _(tensor: torch.Tensor) -> torch.Tensor:
+    return tensor.clone()
+
+
+@clone.register
+def _(tensors: list) -> List[torch.Tensor]:
+    return list(map(clone, tensors))
+
+
+@clone.register
+def _(tensors: dict) -> Dict[str, torch.Tensor]:
+    return {key: clone(tensors[key]) for key in tensors}

eva/core/utils/memory.py

@@ -0,0 +1,28 @@
+"""Memory related functions."""
+
+import functools
+from typing import Any, Dict, List
+
+import torch
+
+
+@functools.singledispatch
+def to_cpu(tensor_type: Any) -> Any:
+    """Moves tensor objects to `cpu`."""
+    raise TypeError(f"Unsupported input type: {type(tensor_type)}.")
+
+
+@to_cpu.register
+def _(tensor: torch.Tensor) -> torch.Tensor:
+    detached_tensor = tensor.detach()
+    return detached_tensor.cpu()
+
+
+@to_cpu.register
+def _(tensors: list) -> List[torch.Tensor]:
+    return list(map(to_cpu, tensors))
+
+
+@to_cpu.register
+def _(tensors: dict) -> Dict[str, torch.Tensor]:
+    return {key: to_cpu(tensors[key]) for key in tensors}

eva/core/utils/operations.py

@@ -0,0 +1,26 @@
+"""Functional operations."""
+
+import re
+from typing import Iterable, List
+
+
+def numeric_sort(item: Iterable[str], /) -> List[str]:
+    """Sorts an iterable of strings treating embedded numbers as numeric values.
+
+    Here the strings are compared based on their numeric value rather than their
+    string representation.
+
+    Args:
+        item: An iterable of strings to be sorted.
+
+    Returns:
+        A list of strings sorted based on their numeric values.
+    """
+    return sorted(
+        item,
+        key=lambda value: re.sub(
+            r"(\d+)",
+            lambda num: f"{int(num[0]):010d}",
+            value,
+        ),
+    )

eva/core/utils/parser.py

@@ -0,0 +1,20 @@
+"""Parsing related helper functions."""
+
+from typing import Any, Dict
+
+import jsonargparse
+
+
+def parse_object(config: Dict[str, Any], expected_type: Any = Any) -> Any:
+    """Parse object which is defined as dictionary."""
+    parser = jsonargparse.ArgumentParser()
+    parser.add_argument("module", type=expected_type)
+    configuration = parser.parse_object({"module": config})
+    init_object = parser.instantiate_classes(configuration)
+    obj_module = init_object.module
+    if isinstance(obj_module, jsonargparse.Namespace):
+        raise ValueError(
+            f"Failed to parsed object '{obj_module.class_path}'. "
+            "Please check that the initialized arguments are valid."
+        )
+    return obj_module

eva/vision/__init__.py

@@ -1,7 +1,7 @@
 """eva vision API."""
 
 try:
-    from eva.vision import models, utils
+    from eva.vision import callbacks, losses, models, utils
     from eva.vision.data import datasets, transforms
 except ImportError as e:
     msg = (
@@ -11,4 +11,4 @@ except ImportError as e:
     )
     raise ImportError(str(e) + "\n\n" + msg) from e
 
-__all__ = ["models", "utils", "datasets", "transforms"]
+__all__ = ["callbacks", "losses", "models", "utils", "datasets", "transforms"]

eva/vision/callbacks/__init__.py

@@ -0,0 +1,5 @@
+"""Vision callbacks API."""
+
+from eva.vision.callbacks.loggers import SemanticSegmentationLogger
+
+__all__ = ["SemanticSegmentationLogger"]

eva/vision/callbacks/loggers/__init__.py

@@ -0,0 +1,5 @@
+"""Vision logging related callbacks API."""
+
+from eva.vision.callbacks.loggers.batch import SemanticSegmentationLogger
+
+__all__ = ["SemanticSegmentationLogger"]

eva/vision/callbacks/loggers/batch/__init__.py

@@ -0,0 +1,5 @@
+"""Batch related loggers callbacks API."""
+
+from eva.vision.callbacks.loggers.batch.segmentation import SemanticSegmentationLogger
+
+__all__ = ["SemanticSegmentationLogger"]