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.

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
save_head_only	bool	Whether to save only the head during checkpointing. If False, will also save the backbone (not recommended when frozen).	True

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,
    save_head_only: bool = True,
) -> 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.
        save_head_only: Whether to save only the head during checkpointing. If False,
            will also save the backbone (not recommended when frozen).
    """
    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
    self.save_head_only = save_head_only

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