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

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

@@ -0,0 +1,76 @@
+"""Default metric collection for binary classification tasks."""
+
+from torchmetrics import classification
+
+from eva.core.metrics import binary_balanced_accuracy, structs
+
+
+class BinaryClassificationMetrics(structs.MetricCollection):
+    """Default metrics for binary classification tasks."""
+
+    def __init__(
+        self,
+        threshold: float = 0.5,
+        ignore_index: int | None = None,
+        prefix: str | None = None,
+        postfix: str | None = None,
+    ) -> None:
+        """Initializes the binary classification metrics.
+
+        The metrics instantiated here are:
+
+        - BinaryAUROC
+        - BinaryAccuracy
+        - BinaryBalancedAccuracy
+        - BinaryF1Score
+        - BinaryPrecision
+        - BinaryRecall
+
+        Args:
+            threshold: Threshold for transforming probability to binary (0,1) predictions
+            ignore_index: Specifies a target value that is ignored and does not
+                contribute to the metric calculation.
+            prefix: A string to append in front of the keys of the output dict.
+            postfix: A string to append after the keys of the output dict.
+        """
+        super().__init__(
+            metrics=[
+                classification.BinaryAUROC(
+                    ignore_index=ignore_index,
+                ),
+                classification.BinaryAccuracy(
+                    threshold=threshold,
+                    ignore_index=ignore_index,
+                ),
+                binary_balanced_accuracy.BinaryBalancedAccuracy(
+                    threshold=threshold,
+                    ignore_index=ignore_index,
+                ),
+                classification.BinaryF1Score(
+                    threshold=threshold,
+                    ignore_index=ignore_index,
+                ),
+                classification.BinaryPrecision(
+                    threshold=threshold,
+                    ignore_index=ignore_index,
+                ),
+                classification.BinaryRecall(
+                    threshold=threshold,
+                    ignore_index=ignore_index,
+                ),
+            ],
+            prefix=prefix,
+            postfix=postfix,
+            compute_groups=[
+                [
+                    "BinaryAccuracy",
+                    "BinaryBalancedAccuracy",
+                    "BinaryF1Score",
+                    "BinaryPrecision",
+                    "BinaryRecall",
+                ],
+                [
+                    "BinaryAUROC",
+                ],
+            ],
+        )

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

@@ -0,0 +1,80 @@
+"""Default metric collection for multiclass classification tasks."""
+
+from typing import Literal
+
+from torchmetrics import classification
+
+from eva.core.metrics import structs
+
+
+class MulticlassClassificationMetrics(structs.MetricCollection):
+    """Default metrics for multi-class classification tasks."""
+
+    def __init__(
+        self,
+        num_classes: int,
+        average: Literal["macro", "weighted", "none"] = "macro",
+        ignore_index: int | None = None,
+        prefix: str | None = None,
+        postfix: str | None = None,
+    ) -> None:
+        """Initializes the multi-class classification metrics.
+
+        The metrics instantiated here are:
+
+        - MulticlassAccuracy
+        - MulticlassPrecision
+        - MulticlassRecall
+        - MulticlassF1Score
+        - MulticlassAUROC
+
+        Args:
+            num_classes: Integer specifying the number of classes.
+            average: Defines the reduction that is applied over labels.
+            ignore_index: Specifies a target value that is ignored and does not
+                contribute to the metric calculation.
+            prefix: A string to append in front of the keys of the output dict.
+            postfix: A string to append after the keys of the output dict.
+        """
+        super().__init__(
+            metrics=[
+                classification.MulticlassAUROC(
+                    num_classes=num_classes,
+                    average=average,
+                    ignore_index=ignore_index,
+                ),
+                classification.MulticlassAccuracy(
+                    num_classes=num_classes,
+                    average=average,
+                    ignore_index=ignore_index,
+                ),
+                classification.MulticlassF1Score(
+                    num_classes=num_classes,
+                    average=average,
+                    ignore_index=ignore_index,
+                ),
+                classification.MulticlassPrecision(
+                    num_classes=num_classes,
+                    average=average,
+                    ignore_index=ignore_index,
+                ),
+                classification.MulticlassRecall(
+                    num_classes=num_classes,
+                    average=average,
+                    ignore_index=ignore_index,
+                ),
+            ],
+            prefix=prefix,
+            postfix=postfix,
+            compute_groups=[
+                [
+                    "MulticlassAccuracy",
+                    "MulticlassF1Score",
+                    "MulticlassPrecision",
+                    "MulticlassRecall",
+                ],
+                [
+                    "MulticlassAUROC",
+                ],
+            ],
+        )

eva/core/metrics/structs/__init__.py

@@ -0,0 +1,9 @@
+"""Core metrics modules API."""
+
+from eva.core.metrics.structs.collection import MetricCollection
+from eva.core.metrics.structs.metric import Metric
+from eva.core.metrics.structs.module import MetricModule
+from eva.core.metrics.structs.schemas import MetricsSchema
+from eva.core.metrics.structs.typings import MetricModuleType
+
+__all__ = ["MetricCollection", "Metric", "MetricModule", "MetricsSchema", "MetricModuleType"]

eva/core/metrics/structs/collection.py

@@ -0,0 +1,6 @@
+"""Metric collection aggregator."""
+
+import torchmetrics
+
+MetricCollection = torchmetrics.MetricCollection
+"""Defines a metric aggregator object."""

eva/core/metrics/structs/metric.py

@@ -0,0 +1,6 @@
+"""Base class of metrics."""
+
+import torchmetrics
+
+Metric = torchmetrics.Metric
+"""Abstract metric class."""

eva/core/metrics/structs/module.py

@@ -0,0 +1,115 @@
+"""Metrics module."""
+
+from __future__ import annotations
+
+from torch import nn
+
+from eva.core.metrics.structs import collection, schemas
+from eva.core.metrics.structs.typings import MetricModuleType
+
+
+class MetricModule(nn.Module):
+    """The metrics module.
+
+    Allows to store and keep track of `train`, `val` and `test` metrics.
+    """
+
+    def __init__(
+        self,
+        train: collection.MetricCollection | None,
+        val: collection.MetricCollection | None,
+        test: collection.MetricCollection | None,
+    ) -> None:
+        """Initializes the metrics for the Trainer.
+
+        Args:
+            train: The training metric collection.
+            val: The validation metric collection.
+            test: The test metric collection.
+        """
+        super().__init__()
+
+        self._train = train or self.default_metric_collection
+        self._val = val or self.default_metric_collection
+        self._test = test or self.default_metric_collection
+
+    @property
+    def default_metric_collection(self) -> collection.MetricCollection:
+        """Returns the default metric collection."""
+        return collection.MetricCollection([])
+
+    @classmethod
+    def from_metrics(
+        cls,
+        train: MetricModuleType | None,
+        val: MetricModuleType | None,
+        test: MetricModuleType | None,
+        *,
+        separator: str = "/",
+    ) -> MetricModule:
+        """Initializes a metric module from a list of metrics.
+
+        Args:
+            train: Metrics for the training stage.
+            val: Metrics for the validation stage.
+            test: Metrics for the test stage.
+            separator: The separator between the group name of the metric
+                and the metric itself.
+        """
+        return cls(
+            train=_create_collection_from_metrics(train, prefix="train" + separator),
+            val=_create_collection_from_metrics(val, prefix="val" + separator),
+            test=_create_collection_from_metrics(test, prefix="test" + separator),
+        )
+
+    @classmethod
+    def from_schema(
+        cls,
+        schema: schemas.MetricsSchema,
+        *,
+        separator: str = "/",
+    ) -> MetricModule:
+        """Initializes a metric module from the metrics schema.
+
+        Args:
+            schema: The dataclass metric schema.
+            separator: The separator between the group name of the metric
+                and the metric itself.
+        """
+        return cls.from_metrics(
+            train=schema.training_metrics,
+            val=schema.evaluation_metrics,
+            test=schema.evaluation_metrics,
+            separator=separator,
+        )
+
+    @property
+    def training_metrics(self) -> collection.MetricCollection:
+        """Returns the metrics of the train dataset."""
+        return self._train
+
+    @property
+    def validation_metrics(self) -> collection.MetricCollection:
+        """Returns the metrics of the validation dataset."""
+        return self._val
+
+    @property
+    def test_metrics(self) -> collection.MetricCollection:
+        """Returns the metrics of the test dataset."""
+        return self._test
+
+
+def _create_collection_from_metrics(
+    metrics: MetricModuleType | None, *, prefix: str | None = None
+) -> collection.MetricCollection | None:
+    """Create a unique collection from metrics.
+
+    Args:
+        metrics: The desired metrics.
+        prefix: A prefix to added to the collection.
+
+    Returns:
+        The resulted metrics collection.
+    """
+    metrics_collection = collection.MetricCollection(metrics or [], prefix=prefix)  # type: ignore
+    return metrics_collection.clone()

eva/core/metrics/structs/schemas.py

@@ -0,0 +1,47 @@
+"""Metrics related helper schemas."""
+
+import dataclasses
+
+from eva.core.metrics.structs.typings import MetricModuleType
+
+
+@dataclasses.dataclass(frozen=True)
+class MetricsSchema:
+    """Metrics schema."""
+
+    common: MetricModuleType | None = None
+    """Holds the common train and evaluation metrics."""
+
+    train: MetricModuleType | None = None
+    """The exclusive training metrics."""
+
+    evaluation: MetricModuleType | None = None
+    """The exclusive evaluation metrics."""
+
+    @property
+    def training_metrics(self) -> MetricModuleType | None:
+        """Returns the training metics."""
+        return self._join_with_common(self.train)
+
+    @property
+    def evaluation_metrics(self) -> MetricModuleType | None:
+        """Returns the evaluation metics."""
+        return self._join_with_common(self.evaluation)
+
+    def _join_with_common(self, metrics: MetricModuleType | None) -> MetricModuleType | None:
+        """Joins the provided metrics with the common.
+
+        Note that if there is duplication of metrics between the provided and the common
+        (meaning there is the same metric in `metrics` and in `self.common`) both will
+        be preserved.
+
+        Args:
+            metrics: The metrics to join.
+
+        Returns:
+            The resulted metrics after joining with the common ones.
+        """
+        if metrics is None or self.common is None:
+            return self.common or metrics
+
+        return [self.common, metrics]  # type: ignore

eva/core/metrics/structs/typings.py

@@ -0,0 +1,15 @@
+"""Metric typings."""
+
+from typing import Dict, Sequence, Union
+
+from eva.core.metrics.structs import collection, metric
+
+BaseMetricModuleType = Union[metric.Metric, collection.MetricCollection]
+"""The base module metric type."""
+
+MetricModuleType = Union[
+    BaseMetricModuleType,
+    Sequence[BaseMetricModuleType],
+    Dict[str, BaseMetricModuleType],
+]
+"""The module metric type."""

eva/core/models/__init__.py

@@ -0,0 +1,13 @@
+"""Models API."""
+
+from eva.core.models.modules import HeadModule, InferenceModule
+from eva.core.models.networks import MLP, HuggingFaceModel, ModelFromFunction, ONNXModel
+
+__all__ = [
+    "HeadModule",
+    "InferenceModule",
+    "MLP",
+    "HuggingFaceModel",
+    "ModelFromFunction",
+    "ONNXModel",
+]

eva/core/models/modules/__init__.py

@@ -0,0 +1,7 @@
+"""Model Modules API."""
+
+from eva.core.models.modules.head import HeadModule
+from eva.core.models.modules.inference import InferenceModule
+from eva.core.models.modules.module import ModelModule
+
+__all__ = ["HeadModule", "ModelModule", "InferenceModule"]

eva/core/models/modules/head.py

@@ -0,0 +1,113 @@
+""""Neural Network Head Module."""
+
+from typing import Any, Callable
+
+import torch
+from lightning.pytorch.cli import LRSchedulerCallable, OptimizerCallable
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+from torch import optim
+from torch.optim import lr_scheduler
+from typing_extensions import override
+
+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
+
+
+class HeadModule(module.ModelModule):
+    """Neural Net Head Module for training on features.
+
+    It can be used for supervised (mini-batch) stochastic gradient descent
+    downstream tasks such as classification, regression and segmentation.
+    """
+
+    def __init__(
+        self,
+        head: MODEL_TYPE,
+        criterion: Callable[..., torch.Tensor],
+        backbone: MODEL_TYPE | None = None,
+        optimizer: OptimizerCallable = optim.Adam,
+        lr_scheduler: LRSchedulerCallable = lr_scheduler.ConstantLR,
+        metrics: metrics_lib.MetricsSchema | None = None,
+        postprocess: batch_postprocess.BatchPostProcess | None = None,
+    ) -> None:
+        """Initializes the neural net head module.
+
+        Args:
+            head: The neural network that would be trained on the features.
+            criterion: The loss function to use.
+            backbone: The feature extractor. If `None`, it will be expected
+                that the input batch returns the features directly.
+            optimizer: The optimizer to use.
+            lr_scheduler: The learning rate scheduler to use.
+            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__(metrics=metrics, postprocess=postprocess)
+
+        self.head = head
+        self.criterion = criterion
+        self.backbone = backbone
+        self.optimizer = optimizer
+        self.lr_scheduler = lr_scheduler
+
+    @override
+    def configure_optimizers(self) -> Any:
+        parameters = list(self.head.parameters())
+        optimizer = self.optimizer(parameters)
+        lr_scheduler = self.lr_scheduler(optimizer)
+        return {"optimizer": optimizer, "lr_scheduler": lr_scheduler}
+
+    @override
+    def forward(self, tensor: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
+        features = tensor if self.backbone is None else self.backbone(tensor)
+        return self.head(features).squeeze(-1)
+
+    @override
+    def on_fit_start(self) -> None:
+        if self.backbone is not None:
+            grad.deactivate_requires_grad(self.backbone)
+
+    @override
+    def training_step(self, batch: INPUT_BATCH, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
+        return self._batch_step(batch)
+
+    @override
+    def validation_step(self, batch: INPUT_BATCH, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
+        return self._batch_step(batch)
+
+    @override
+    def test_step(self, batch: INPUT_BATCH, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
+        return self._batch_step(batch)
+
+    @override
+    def predict_step(self, batch: INPUT_BATCH, *args: Any, **kwargs: Any) -> torch.Tensor:
+        tensor = INPUT_BATCH(*batch).data
+        return tensor if self.backbone is None else self.backbone(tensor)
+
+    @override
+    def on_fit_end(self) -> None:
+        if self.backbone is not None:
+            grad.activate_requires_grad(self.backbone)
+
+    def _batch_step(self, batch: INPUT_BATCH) -> STEP_OUTPUT:
+        """Performs a model forward step and calculates the loss.
+
+        Args:
+            batch: The desired batch to process.
+
+        Returns:
+            The batch step output.
+        """
+        data, targets, metadata = INPUT_BATCH(*batch)
+        predictions = self(data)
+        loss = self.criterion(predictions, targets)
+        return {
+            "loss": loss,
+            "targets": targets,
+            "predictions": predictions,
+            "metadata": metadata,
+        }

eva/core/models/modules/inference.py

@@ -0,0 +1,37 @@
+"""Model inference module."""
+
+import torch
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+from typing_extensions import override
+
+from eva.core.models.modules import module
+from eva.core.models.modules.typings import INPUT_BATCH, MODEL_TYPE
+
+
+class InferenceModule(module.ModelModule):
+    """An lightweight model module to perform inference."""
+
+    def __init__(self, backbone: MODEL_TYPE) -> None:
+        """Initializes the module.
+
+        Args:
+            backbone: The network to be used for inference.
+        """
+        super().__init__(metrics=None)
+
+        self.backbone = backbone
+
+    @override
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        return self.backbone(tensor)
+
+    @override
+    def predict_step(
+        self,
+        batch: INPUT_BATCH,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> STEP_OUTPUT:
+        data, *_ = INPUT_BATCH(*batch)
+        predictions = self(data)
+        return predictions