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

eva/core/models/modules/module.py

@@ -4,6 +4,7 @@ from typing import Any, Mapping
 
 import lightning.pytorch as pl
 import torch
+from lightning.pytorch.strategies.single_device import SingleDeviceStrategy
 from lightning.pytorch.utilities import memory
 from lightning.pytorch.utilities.types import STEP_OUTPUT
 from typing_extensions import override
@@ -46,6 +47,21 @@ class ModelModule(pl.LightningModule):
         """The default post-processes."""
         return batch_postprocess.BatchPostProcess()
 
+    @property
+    def metrics_device(self) -> torch.device:
+        """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.
+        """
+        move_to_cpu = isinstance(self.trainer.strategy, SingleDeviceStrategy)
+        return torch.device("cpu") if move_to_cpu else self.device
+
+    @override
+    def on_fit_start(self) -> None:
+        self.metrics.to(device=self.metrics_device)
+
     @override
     def on_train_batch_end(
         self,
@@ -59,6 +75,10 @@ class ModelModule(pl.LightningModule):
             batch_outputs=outputs,
         )
 
+    @override
+    def on_validation_start(self) -> None:
+        self.metrics.to(device=self.metrics_device)
+
     @override
     def on_validation_batch_end(
         self,
@@ -78,6 +98,10 @@ class ModelModule(pl.LightningModule):
     def on_validation_epoch_end(self) -> None:
         self._compute_and_log_metrics(self.metrics.validation_metrics)
 
+    @override
+    def on_test_start(self) -> None:
+        self.metrics.to(device=self.metrics_device)
+
     @override
     def on_test_batch_end(
         self,
@@ -110,7 +134,7 @@ class ModelModule(pl.LightningModule):
             The updated outputs.
         """
         self._postprocess(outputs)
-        return memory.recursive_detach(outputs, to_cpu=self.device.type == "cpu")
+        return memory.recursive_detach(outputs, to_cpu=self.metrics_device.type == "cpu")
 
     def _forward_and_log_metrics(
         self,

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/_recorder.py

@@ -5,18 +5,41 @@ import json
 import os
 import statistics
 import sys
-from typing import Any, Dict, List, Mapping
+from typing import Dict, List, Mapping, TypedDict
 
 from lightning.pytorch.utilities.types import _EVALUATE_OUTPUT
 from lightning_fabric.utilities import cloud_io
 from loguru import logger
 from omegaconf import OmegaConf
+from rich import console as rich_console
+from rich import table as rich_table
 from toolz import dicttoolz
 
 SESSION_METRICS = Mapping[str, List[float]]
 """Session metrics type-hint."""
 
 
+class SESSION_STATISTICS(TypedDict):
+    """Type-hint for aggregated metrics of multiple runs with mean & stdev."""
+
+    mean: float
+    stdev: float
+    values: List[float]
+
+
+class STAGE_RESULTS(TypedDict):
+    """Type-hint for metrics statstics for val & test stages."""
+
+    val: List[Dict[str, SESSION_STATISTICS]]
+    test: List[Dict[str, SESSION_STATISTICS]]
+
+
+class RESULTS_DICT(TypedDict):
+    """Type-hint for the final results dictionary."""
+
+    metrics: STAGE_RESULTS
+
+
 class SessionRecorder:
     """Multi-run (session) summary logger."""
 
@@ -25,6 +48,7 @@ class SessionRecorder:
         output_dir: str,
         results_file: str = "results.json",
         config_file: str = "config.yaml",
+        verbose: bool = True,
     ) -> None:
         """Initializes the recorder.
 
@@ -32,10 +56,12 @@ class SessionRecorder:
             output_dir: The destination folder to save the results.
             results_file: The name of the results json file.
             config_file: The name of the yaml configuration file.
+            verbose: Whether to print the session metrics.
         """
         self._output_dir = output_dir
         self._results_file = results_file
         self._config_file = config_file
+        self._verbose = verbose
 
         self._validation_metrics: List[SESSION_METRICS] = []
         self._test_metrics: List[SESSION_METRICS] = []
@@ -67,13 +93,13 @@ class SessionRecorder:
         self._update_validation_metrics(validation_scores)
         self._update_test_metrics(test_scores)
 
-    def compute(self) -> Dict[str, List[Dict[str, Any]]]:
+    def compute(self) -> STAGE_RESULTS:
         """Computes and returns the session statistics."""
         validation_statistics = list(map(_calculate_statistics, self._validation_metrics))
         test_statistics = list(map(_calculate_statistics, self._test_metrics))
         return {"val": validation_statistics, "test": test_statistics}
 
-    def export(self) -> Dict[str, Any]:
+    def export(self) -> RESULTS_DICT:
         """Exports the results."""
         statistics = self.compute()
         return {"metrics": statistics}
@@ -83,6 +109,8 @@ class SessionRecorder:
         results = self.export()
         _save_json(results, self.filename)
         self._save_config()
+        if self._verbose:
+            _print_results(results)
 
     def reset(self) -> None:
         """Resets the state of the tracked metrics."""
@@ -125,10 +153,10 @@ def _init_session_metrics(n_datasets: int) -> List[SESSION_METRICS]:
     return [collections.defaultdict(list) for _ in range(n_datasets)]
 
 
-def _calculate_statistics(session_metrics: SESSION_METRICS) -> Dict[str, float | List[float]]:
+def _calculate_statistics(session_metrics: SESSION_METRICS) -> Dict[str, SESSION_STATISTICS]:
     """Calculate the metric statistics of a dataset session run."""
 
-    def _calculate_metric_statistics(values: List[float]) -> Dict[str, float | List[float]]:
+    def _calculate_metric_statistics(values: List[float]) -> SESSION_STATISTICS:
         """Calculates and returns the metric statistics."""
         mean = statistics.mean(values)
         stdev = statistics.stdev(values) if len(values) > 1 else 0
@@ -137,7 +165,7 @@ def _calculate_statistics(session_metrics: SESSION_METRICS) -> Dict[str, float |
     return dicttoolz.valmap(_calculate_metric_statistics, session_metrics)
 
 
-def _save_json(data: Dict[str, Any], save_as: str = "data.json"):
+def _save_json(data: RESULTS_DICT, save_as: str = "data.json"):
     """Saves data to a json file."""
     if not save_as.endswith(".json"):
         raise ValueError()
@@ -146,4 +174,38 @@ def _save_json(data: Dict[str, Any], save_as: str = "data.json"):
     fs = cloud_io.get_filesystem(output_dir, anon=False)
     fs.makedirs(output_dir, exist_ok=True)
     with fs.open(save_as, "w") as file:
-        json.dump(data, file, indent=4, sort_keys=True)
+        json.dump(data, file, indent=2, sort_keys=True)
+
+
+def _print_results(results: RESULTS_DICT) -> None:
+    """Prints the results to the console."""
+    try:
+        for stage in ["val", "test"]:
+            for dataset_idx in range(len(results["metrics"][stage])):
+                _print_table(results["metrics"][stage][dataset_idx], stage, dataset_idx)
+    except Exception as e:
+        logger.error(f"Failed to print the results: {e}")
+
+
+def _print_table(metrics_dict: Dict[str, SESSION_STATISTICS], stage: str, dataset_idx: int):
+    """Prints the metrics of a single dataset as a table."""
+    metrics_table = rich_table.Table(
+        title=f"\n{stage.capitalize()} Dataset {dataset_idx}", title_style="bold"
+    )
+    metrics_table.add_column("Metric", style="cyan")
+    metrics_table.add_column("Mean", style="magenta")
+    metrics_table.add_column("Stdev", style="magenta")
+    metrics_table.add_column("All", style="magenta")
+
+    n_runs = len(metrics_dict[next(iter(metrics_dict))]["values"])
+    for metric_name, metric_dict in metrics_dict.items():
+        row = [
+            metric_name,
+            f'{metric_dict["mean"]:.3f}',
+            f'{metric_dict["stdev"]:.3f}',
+            ", ".join(f'{metric_dict["values"][i]:.3f}' for i in range(n_runs)),
+        ]
+        metrics_table.add_row(*row)
+
+    console = rich_console.Console()
+    console.print(metrics_table)

eva/core/trainers/functional.py

@@ -16,6 +16,7 @@ def run_evaluation_session(
     datamodule: datamodules.DataModule,
     *,
     n_runs: int = 1,
+    verbose: bool = True,
 ) -> None:
     """Runs a downstream evaluation session out-of-place.
 
@@ -29,11 +30,17 @@ def run_evaluation_session(
         base_model: The base model module to use.
         datamodule: The data module.
         n_runs: The amount of runs (fit and evaluate) to perform.
+        verbose: Whether to verbose the session metrics instead of
+            these of each individual runs and vice-versa.
     """
-    recorder = _recorder.SessionRecorder(output_dir=base_trainer.default_log_dir)
+    recorder = _recorder.SessionRecorder(output_dir=base_trainer.default_log_dir, verbose=verbose)
     for run_index in range(n_runs):
         validation_scores, test_scores = run_evaluation(
-            base_trainer, base_model, datamodule, run_id=f"run_{run_index}"
+            base_trainer,
+            base_model,
+            datamodule,
+            run_id=f"run_{run_index}",
+            verbose=not verbose,
         )
         recorder.update(validation_scores, test_scores)
     recorder.save()
@@ -45,6 +52,7 @@ def run_evaluation(
     datamodule: datamodules.DataModule,
     *,
     run_id: str | None = None,
+    verbose: bool = True,
 ) -> Tuple[_EVALUATE_OUTPUT, _EVALUATE_OUTPUT | None]:
     """Fits and evaluates a model out-of-place.
 
@@ -54,19 +62,23 @@ def run_evaluation(
         datamodule: The data module.
         run_id: The run id to be appended to the output log directory.
             If `None`, it will use the log directory of the trainer as is.
+        verbose: Whether to print the validation and test metrics
+            in the end of the training.
 
     Returns:
         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)
+    return fit_and_validate(trainer, model, datamodule, verbose=verbose)
 
 
 def fit_and_validate(
     trainer: eva_trainer.Trainer,
     model: modules.ModelModule,
     datamodule: datamodules.DataModule,
+    verbose: bool = True,
 ) -> Tuple[_EVALUATE_OUTPUT, _EVALUATE_OUTPUT | None]:
     """Fits and evaluates a model in-place.
 
@@ -77,13 +89,19 @@ def fit_and_validate(
         trainer: The trainer module to use and update in-place.
         model: The model module to use and update in-place.
         datamodule: The data module.
+        verbose: Whether to print the validation and test metrics
+            in the end of the training.
 
     Returns:
         A tuple of with the validation and the test metrics (if exists).
     """
     trainer.fit(model, datamodule=datamodule)
-    validation_scores = trainer.validate(datamodule=datamodule)
-    test_scores = None if datamodule.datasets.test is None else trainer.test(datamodule=datamodule)
+    validation_scores = trainer.validate(datamodule=datamodule, verbose=verbose)
+    test_scores = (
+        None
+        if datamodule.datasets.test is None
+        else trainer.test(datamodule=datamodule, verbose=verbose)
+    )
     return validation_scores, test_scores