Modules

Reference information for the model Modules API.

eva.models.modules.ModelModule

Bases: LightningModule

The base model module.

Parameters:

Name	Type	Description	Default
metrics	MetricsSchema | None	The metric groups to track.	None
postprocess	BatchPostProcess | None	A list of helper functions to apply after the loss and before the metrics calculation to the model predictions and targets.	None

Source code in src/eva/core/models/modules/module.py

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)

default_metrics: metrics_lib.MetricsSchema property

The default metrics.

default_postprocess: batch_postprocess.BatchPostProcess property

The default post-processes.

metrics_device: torch.device property

Returns the device by which the metrics should be calculated.

We allocate the metrics to CPU when operating on single device, as it is much faster, but to GPU when employing multiple ones, as DDP strategy requires the metrics to be allocated to the module's GPU.

eva.models.modules.HeadModule

Bases: 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.

Parameters:

Name	Type	Description	Default
head	Dict[str, Any] | MODEL_TYPE	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.	required
criterion	Callable[..., Tensor]	The loss function to use.	required
backbone	MODEL_TYPE | None	The feature extractor. If None, it will be expected that the input batch returns the features directly.	None
optimizer	OptimizerCallable	The optimizer to use.	Adam
lr_scheduler	LRSchedulerCallable	The learning rate scheduler to use.	ConstantLR
metrics	MetricsSchema | None	The metric groups to track.	None
postprocess	BatchPostProcess | None	A list of helper functions to apply after the loss and before the metrics calculation to the model predictions and targets.	None

Source code in src/eva/core/models/modules/head.py

def __init__(
    self,
    head: Dict[str, Any] | 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.
            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.
        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  # type: ignore
    self.criterion = criterion
    self.backbone = backbone
    self.optimizer = optimizer
    self.lr_scheduler = lr_scheduler

eva.models.modules.InferenceModule

Bases: ModelModule

An lightweight model module to perform inference.

Parameters:

Name	Type	Description	Default
backbone	MODEL_TYPE	The network to be used for inference.	required

Source code in src/eva/core/models/modules/inference.py

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