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

eva/core/data/datamodules/datamodule.py

@@ -0,0 +1,108 @@
+"""Core DataModule."""
+
+from typing import List
+
+import lightning.pytorch as pl
+from lightning.pytorch.utilities.types import EVAL_DATALOADERS, TRAIN_DATALOADERS
+from typing_extensions import override
+
+from eva.core.data import dataloaders as dataloaders_lib
+from eva.core.data import datasets as datasets_lib
+from eva.core.data.datamodules import call, schemas
+
+
+class DataModule(pl.LightningDataModule):
+    """DataModule encapsulates all the steps needed to process data.
+
+    It will initialize and create the mapping between dataloaders and
+    datasets. During the `prepare_data`, `setup` and `teardown`, the
+    datamodule will call the respective methods from all datasets,
+    given that they are defined.
+    """
+
+    def __init__(
+        self,
+        datasets: schemas.DatasetsSchema | None = None,
+        dataloaders: schemas.DataloadersSchema | None = None,
+    ) -> None:
+        """Initializes the datamodule.
+
+        Args:
+            datasets: The desired datasets.
+            dataloaders: The desired dataloaders.
+        """
+        super().__init__()
+
+        self.datasets = datasets or self.default_datasets
+        self.dataloaders = dataloaders or self.default_dataloaders
+
+    @property
+    def default_datasets(self) -> schemas.DatasetsSchema:
+        """Returns the default datasets."""
+        return schemas.DatasetsSchema()
+
+    @property
+    def default_dataloaders(self) -> schemas.DataloadersSchema:
+        """Returns the default dataloader schema."""
+        return schemas.DataloadersSchema()
+
+    @override
+    def prepare_data(self) -> None:
+        call.call_method_if_exists(self.datasets.tolist(), "prepare_data")
+
+    @override
+    def setup(self, stage: str) -> None:
+        call.call_method_if_exists(self.datasets.tolist(stage), "setup")
+
+    @override
+    def teardown(self, stage: str) -> None:
+        call.call_method_if_exists(self.datasets.tolist(stage), "teardown")
+
+    @override
+    def train_dataloader(self) -> TRAIN_DATALOADERS:
+        if self.datasets.train is None:
+            raise ValueError(
+                "Train dataloader can not be initialized as `self.datasets.train` is `None`."
+            )
+        return self.dataloaders.train(self.datasets.train)
+
+    @override
+    def val_dataloader(self) -> EVAL_DATALOADERS:
+        if self.datasets.val is None:
+            raise ValueError(
+                "Validation dataloader can not be initialized as `self.datasets.val` is `None`."
+            )
+        return self._initialize_dataloaders(self.dataloaders.val, self.datasets.val)
+
+    @override
+    def test_dataloader(self) -> EVAL_DATALOADERS:
+        if self.datasets.test is None:
+            raise ValueError(
+                "Test dataloader can not be initialized as `self.datasets.test` is `None`."
+            )
+        return self._initialize_dataloaders(self.dataloaders.test, self.datasets.test)
+
+    @override
+    def predict_dataloader(self) -> EVAL_DATALOADERS:
+        if self.datasets.predict is None:
+            raise ValueError(
+                "Predict dataloader can not be initialized as `self.datasets.predict` is `None`."
+            )
+        return self._initialize_dataloaders(self.dataloaders.predict, self.datasets.predict)
+
+    def _initialize_dataloaders(
+        self,
+        dataloader: dataloaders_lib.DataLoader,
+        datasets: datasets_lib.TorchDataset | List[datasets_lib.TorchDataset],
+    ) -> EVAL_DATALOADERS:
+        """Initializes dataloaders from a given set of dataset.
+
+        Args:
+            dataloader: The dataloader to apply to the provided datasets.
+            datasets: The desired dataset(s) to allocate dataloader(s).
+
+        Returns:
+            A list with the dataloaders of the provided dataset(s).
+        """
+        datasets = datasets if isinstance(datasets, list) else [datasets]
+        return list(map(dataloader, datasets))

eva/core/data/datamodules/schemas.py

@@ -0,0 +1,62 @@
+"""Argument schemas used in DataModule."""
+
+import dataclasses
+from typing import List
+
+from eva.core.data import dataloaders, datasets
+
+TRAIN_DATASET = datasets.TorchDataset | None
+"""Train dataset."""
+
+EVAL_DATASET = datasets.TorchDataset | List[datasets.TorchDataset] | None
+"""Evaluation dataset."""
+
+
+@dataclasses.dataclass(frozen=True)
+class DatasetsSchema:
+    """Datasets schema used in DataModule."""
+
+    train: TRAIN_DATASET = None
+    """Train dataset."""
+
+    val: EVAL_DATASET = None
+    """Validation dataset."""
+
+    test: EVAL_DATASET = None
+    """Test dataset."""
+
+    predict: EVAL_DATASET = None
+    """Predict dataset."""
+
+    def tolist(self, stage: str | None = None) -> List[EVAL_DATASET]:
+        """Returns the dataclass as a list and optionally filters it given the stage."""
+        match stage:
+            case "fit":
+                return [self.train, self.val]
+            case "validate":
+                return [self.val]
+            case "test":
+                return [self.test]
+            case "predict":
+                return [self.predict]
+            case None:
+                return [self.train, self.val, self.test, self.predict]
+            case _:
+                raise ValueError(f"Invalid stage `{stage}`.")
+
+
+@dataclasses.dataclass(frozen=True)
+class DataloadersSchema:
+    """Dataloaders schema used in DataModule."""
+
+    train: dataloaders.DataLoader = dataclasses.field(default_factory=dataloaders.DataLoader)
+    """Train dataloader."""
+
+    val: dataloaders.DataLoader = dataclasses.field(default_factory=dataloaders.DataLoader)
+    """Validation dataloader."""
+
+    test: dataloaders.DataLoader = dataclasses.field(default_factory=dataloaders.DataLoader)
+    """Test dataloader."""
+
+    predict: dataloaders.DataLoader = dataclasses.field(default_factory=dataloaders.DataLoader)
+    """Predict dataloader."""

eva/core/data/datasets/__init__.py

@@ -0,0 +1,7 @@
+"""Datasets API."""
+
+from eva.core.data.datasets.base import Dataset
+from eva.core.data.datasets.classification import EmbeddingsClassificationDataset
+from eva.core.data.datasets.dataset import TorchDataset
+
+__all__ = ["Dataset", "EmbeddingsClassificationDataset", "TorchDataset"]

eva/core/data/datasets/base.py

@@ -0,0 +1,53 @@
+"""Base dataset class."""
+
+from eva.core.data.datasets import dataset
+
+
+class Dataset(dataset.TorchDataset):
+    """Base dataset class."""
+
+    def prepare_data(self) -> None:
+        """Encapsulates all disk related tasks.
+
+        This method is preferred for downloading and preparing the data, for
+        example generate manifest files. If implemented, it will be called via
+        :class:`eva.core.data.datamodules.DataModule`, which ensures that is called
+        only within a single process, making it multi-processes safe.
+        """
+
+    def setup(self) -> None:
+        """Setups the dataset.
+
+        This method is preferred for creating datasets or performing
+        train/val/test splits. If implemented, it will be called via
+        :class:`eva.core.data.datamodules.DataModule` at the beginning of fit
+        (train + validate), validate, test, or predict and it will be called
+        from every process (i.e. GPU) across all the nodes in DDP.
+        """
+        self.configure()
+        self.validate()
+
+    def configure(self):
+        """Configures the dataset.
+
+        This method is preferred to configure the dataset; assign values
+        to attributes, perform splits etc. This would be called from the
+        method ::method::`setup`, before calling the ::method::`validate`.
+        """
+
+    def validate(self):
+        """Validates the dataset.
+
+        This method aims to check the integrity of the dataset and verify
+        that is configured properly. This would be called from the method
+        ::method::`setup`, after calling the ::method::`configure`.
+        """
+
+    def teardown(self) -> None:
+        """Cleans up the data artifacts.
+
+        Used to clean-up when the run is finished. If implemented, it will
+        be called via :class:`eva.core.data.datamodules.DataModule` at the end
+        of fit (train + validate), validate, test, or predict and it will be
+        called from every process (i.e. GPU) across all the nodes in DDP.
+        """

eva/core/data/datasets/classification/__init__.py

@@ -0,0 +1,5 @@
+"""Classification datasets API."""
+
+from eva.core.data.datasets.classification.embeddings import EmbeddingsClassificationDataset
+
+__all__ = ["EmbeddingsClassificationDataset"]

eva/core/data/datasets/classification/embeddings.py

@@ -0,0 +1,154 @@
+"""Embeddings classification dataset."""
+
+import os
+from typing import Callable, Dict, Tuple
+
+import numpy as np
+import pandas as pd
+import torch
+from typing_extensions import override
+
+from eva.core.data.datasets import base
+from eva.core.utils import io
+
+
+class EmbeddingsClassificationDataset(base.Dataset):
+    """Embeddings classification dataset."""
+
+    default_column_mapping: Dict[str, str] = {
+        "data": "embeddings",
+        "target": "target",
+        "split": "split",
+    }
+    """The default column mapping of the variables to the manifest columns."""
+
+    def __init__(
+        self,
+        root: str,
+        manifest_file: str,
+        split: str | None = None,
+        column_mapping: Dict[str, str] = default_column_mapping,
+        embeddings_transforms: Callable | None = None,
+        target_transforms: Callable | None = None,
+    ) -> None:
+        """Initialize dataset.
+
+        Expects a manifest file listing the paths of .pt files that contain
+        tensor embeddings of shape [embedding_dim] or [1, embedding_dim].
+
+        Args:
+            root: Root directory of the dataset.
+            manifest_file: The path to the manifest file, which is relative to
+                the `root` argument.
+            split: The dataset split to use. The `split` column of the manifest
+                file will be splitted based on this value.
+            column_mapping: Defines the map between the variables and the manifest
+                columns. It will overwrite the `default_column_mapping` with
+                the provided values, so that `column_mapping` can contain only the
+                values which are altered or missing.
+            embeddings_transforms: A function/transform that transforms the embedding.
+            target_transforms: A function/transform that transforms the target.
+        """
+        super().__init__()
+
+        self._root = root
+        self._manifest_file = manifest_file
+        self._split = split
+        self._column_mapping = self.default_column_mapping | column_mapping
+        self._embeddings_transforms = embeddings_transforms
+        self._target_transforms = target_transforms
+
+        self._data: pd.DataFrame
+
+    def filename(self, index: int) -> str:
+        """Returns the filename of the `index`'th data sample.
+
+        Note that this is the relative file path to the root.
+
+        Args:
+            index: The index of the data-sample to select.
+
+        Returns:
+            The filename of the `index`'th data sample.
+        """
+        return self._data.at[index, self._column_mapping["data"]]
+
+    @override
+    def setup(self):
+        self._data = self._load_manifest()
+
+    def __getitem__(self, index) -> Tuple[torch.Tensor, np.ndarray]:
+        """Returns the `index`'th data sample.
+
+        Args:
+            index: The index of the data-sample to select.
+
+        Returns:
+            A data sample and its target.
+        """
+        embeddings = self._load_embeddings(index)
+        target = self._load_target(index)
+        return self._apply_transforms(embeddings, target)
+
+    def __len__(self) -> int:
+        """Returns the total length of the data."""
+        return len(self._data)
+
+    def _load_embeddings(self, index: int) -> torch.Tensor:
+        """Returns the `index`'th embedding sample.
+
+        Args:
+            index: The index of the data sample to load.
+
+        Returns:
+            The sample embedding as an array.
+        """
+        filename = self.filename(index)
+        embeddings_path = os.path.join(self._root, filename)
+        tensor = torch.load(embeddings_path, map_location="cpu")
+        return tensor.squeeze(0)
+
+    def _load_target(self, index: int) -> np.ndarray:
+        """Returns the `index`'th target sample.
+
+        Args:
+            index: The index of the data sample to load.
+
+        Returns:
+            The sample target as an array.
+        """
+        target = self._data.at[index, self._column_mapping["target"]]
+        return np.asarray(target, dtype=np.int64)
+
+    def _load_manifest(self) -> pd.DataFrame:
+        """Loads manifest file and filters the data based on the split column.
+
+        Returns:
+            The data as a pandas DataFrame.
+        """
+        manifest_path = os.path.join(self._root, self._manifest_file)
+        data = io.read_dataframe(manifest_path)
+        if self._split is not None:
+            filtered_data = data.loc[data[self._column_mapping["split"]] == self._split]
+            data = filtered_data.reset_index(drop=True)
+        return data
+
+    def _apply_transforms(
+        self, embeddings: torch.Tensor, target: np.ndarray
+    ) -> Tuple[torch.Tensor, np.ndarray]:
+        """Applies the transforms to the provided data and returns them.
+
+        Args:
+            embeddings: The embeddings to be transformed.
+            target: The training target.
+
+        Returns:
+            A tuple with the embeddings and the target transformed.
+        """
+        if self._embeddings_transforms is not None:
+            embeddings = self._embeddings_transforms(embeddings)
+
+        if self._target_transforms is not None:
+            target = self._target_transforms(target)
+
+        return embeddings, target

eva/core/data/datasets/dataset.py

@@ -0,0 +1,6 @@
+"""Base torch dataset module."""
+
+from torch.utils import data
+
+TorchDataset = data.Dataset
+"""Base torch dataset class."""

eva/core/data/samplers/__init__.py

@@ -0,0 +1,5 @@
+"""Data samplers API."""
+
+from eva.core.data.samplers.sampler import Sampler
+
+__all__ = ["Sampler"]

eva/core/data/samplers/sampler.py

@@ -0,0 +1,6 @@
+"""Core data sampler."""
+
+from torch.utils import data
+
+Sampler = data.Sampler
+"""Core abstract data sampler class."""

eva/core/data/transforms/__init__.py

@@ -0,0 +1,5 @@
+"""Core data transforms."""
+
+from eva.core.data.transforms.dtype import ArrayToFloatTensor, ArrayToTensor
+
+__all__ = ["ArrayToFloatTensor", "ArrayToTensor"]

eva/core/data/transforms/dtype/__init__.py

@@ -0,0 +1,5 @@
+"""Type casting related transforms."""
+
+from eva.core.data.transforms.dtype.array import ArrayToFloatTensor, ArrayToTensor
+
+__all__ = ["ArrayToFloatTensor", "ArrayToTensor"]

eva/core/data/transforms/dtype/array.py

@@ -0,0 +1,28 @@
+"""Transformations to convert numpy arrays to torch tensors."""
+
+import numpy.typing as npt
+import torch
+
+
+class ArrayToTensor:
+    """Converts a numpy array to a torch tensor."""
+
+    def __call__(self, array: npt.ArrayLike) -> torch.Tensor:
+        """Call method for the transformation.
+
+        Args:
+            array: The input numpy array.
+        """
+        return torch.from_numpy(array)
+
+
+class ArrayToFloatTensor(ArrayToTensor):
+    """Converts a numpy array to a torch tensor and casts it to float."""
+
+    def __call__(self, array: npt.ArrayLike):
+        """Call method for the transformation.
+
+        Args:
+            array: The input numpy array.
+        """
+        return super().__call__(array).float()

eva/core/interface/__init__.py

@@ -0,0 +1,5 @@
+"""Interface API."""
+
+from eva.core.interface.interface import Interface
+
+__all__ = ["Interface"]

eva/core/interface/interface.py

@@ -0,0 +1,79 @@
+"""Main interface class."""
+
+from eva.core import trainers as eva_trainer
+from eva.core.data import datamodules
+from eva.core.models import modules
+
+
+class Interface:
+    """A high-level interface for training and validating a machine learning model.
+
+    This class provides a convenient interface to connect a model, data, and trainer
+    to train and validate a model.
+    """
+
+    def fit(
+        self,
+        trainer: eva_trainer.Trainer,
+        model: modules.ModelModule,
+        data: datamodules.DataModule,
+    ) -> None:
+        """Perform model training and evaluation out-of-place.
+
+        This method uses the specified trainer to fit the model using the provided data.
+
+        Example use cases:
+
+        - Using a model consisting of a frozen backbone and a head, the backbone will generate
+          the embeddings on the fly which are then used as input features to train the head on
+          the downstream task specified by the given dataset.
+        - Fitting only the head network using a dataset that loads pre-computed embeddings.
+
+        Args:
+            trainer: The base trainer to use but not modify.
+            model: The model module to use but not modify.
+            data: The data module.
+        """
+        trainer.run_evaluation_session(model=model, datamodule=data)
+
+    def predict(
+        self,
+        trainer: eva_trainer.Trainer,
+        model: modules.ModelModule,
+        data: datamodules.DataModule,
+    ) -> None:
+        """Perform model prediction out-of-place.
+
+        This method performs inference with a pre-trained foundation model to compute embeddings.
+
+        Args:
+            trainer: The base trainer to use but not modify.
+            model: The model module to use but not modify.
+            data: The data module.
+        """
+        eva_trainer.infer_model(
+            base_trainer=trainer,
+            base_model=model,
+            datamodule=data,
+            return_predictions=False,
+        )
+
+    def predict_fit(
+        self,
+        trainer: eva_trainer.Trainer,
+        model: modules.ModelModule,
+        data: datamodules.DataModule,
+    ) -> None:
+        """Combines the predict and fit commands in one method.
+
+        This method performs the following two steps:
+        1. predict: perform inference with a pre-trained foundation model to compute embeddings.
+        2. fit: training the head network using the embeddings generated in step 1.
+
+        Args:
+            trainer: The base trainer to use but not modify.
+            model: The model module to use but not modify.
+            data: The data module.
+        """
+        self.predict(trainer=trainer, model=model, data=data)
+        self.fit(trainer=trainer, model=model, data=data)

eva/core/metrics/__init__.py

@@ -0,0 +1,17 @@
+"""Metrics API."""
+
+from eva.core.metrics.average_loss import AverageLoss
+from eva.core.metrics.binary_balanced_accuracy import BinaryBalancedAccuracy
+from eva.core.metrics.defaults import BinaryClassificationMetrics, MulticlassClassificationMetrics
+from eva.core.metrics.structs import Metric, MetricCollection, MetricModule, MetricsSchema
+
+__all__ = [
+    "AverageLoss",
+    "BinaryBalancedAccuracy",
+    "Metric",
+    "MetricCollection",
+    "MetricModule",
+    "MetricsSchema",
+    "MulticlassClassificationMetrics",
+    "BinaryClassificationMetrics",
+]

eva/core/metrics/average_loss.py

@@ -0,0 +1,47 @@
+"""Implementation of the average loss metric."""
+
+import torch
+from loguru import logger
+from typing_extensions import override
+
+from eva.core.metrics import structs
+
+
+class AverageLoss(structs.Metric):
+    """Average loss metric tracker."""
+
+    is_differentiable = True
+    higher_is_better = False
+    full_state_update = False
+
+    def __init__(self) -> None:
+        """Initializes the metric."""
+        super().__init__()
+
+        self.add_state("value", default=torch.tensor(0), dist_reduce_fx="sum")
+        self.add_state("total", default=torch.tensor(0), dist_reduce_fx="sum")
+
+    @override
+    def update(self, loss: torch.Tensor) -> None:
+        _check_nans(loss)
+        total_samples = loss.numel()
+        if total_samples == 0:
+            return
+
+        self.value = self.value + torch.sum(loss)
+        self.total = self.total + total_samples
+
+    @override
+    def compute(self) -> torch.Tensor:
+        return self.value / self.total
+
+
+def _check_nans(tensor: torch.Tensor) -> None:
+    """Checks for nan values and raises a warning.
+
+    Raises:
+        Warning: If the input tensor consists of any NaN(s).
+    """
+    nan_values = tensor.isnan()
+    if nan_values.any():
+        logger.warning("Encountered `nan` value(s) in input tensor.")

eva/core/metrics/binary_balanced_accuracy.py

@@ -0,0 +1,22 @@
+"""Binary balanced accuracy metric."""
+
+from torch import Tensor
+from torchmetrics.classification import stat_scores
+from torchmetrics.utilities.compute import _safe_divide
+
+
+class BinaryBalancedAccuracy(stat_scores.BinaryStatScores):
+    """Computes the balanced accuracy for binary classification."""
+
+    is_differentiable: bool = False
+    higher_is_better: bool | None = True
+    full_state_update: bool = False
+    plot_lower_bound: float | None = 0.0
+    plot_upper_bound: float | None = 1.0
+
+    def compute(self) -> Tensor:
+        """Compute accuracy based on inputs passed in to ``update`` previously."""
+        tp, fp, tn, fn = self._final_state()
+        sensitivity = _safe_divide(tp, tp + fn)
+        specificity = _safe_divide(tn, tn + fp)
+        return 0.5 * (sensitivity + specificity)

eva/core/metrics/defaults/__init__.py

@@ -0,0 +1,6 @@
+"""Default metric collections API."""
+
+from eva.core.metrics.defaults.classification.binary import BinaryClassificationMetrics
+from eva.core.metrics.defaults.classification.multiclass import MulticlassClassificationMetrics
+
+__all__ = ["MulticlassClassificationMetrics", "BinaryClassificationMetrics"]

eva/core/metrics/defaults/classification/__init__.py

@@ -0,0 +1,6 @@
+"""Default classification metric collections API."""
+
+from eva.core.metrics.defaults.classification.binary import BinaryClassificationMetrics
+from eva.core.metrics.defaults.classification.multiclass import MulticlassClassificationMetrics
+
+__all__ = ["MulticlassClassificationMetrics", "BinaryClassificationMetrics"]