Allow multiple fit runs

configs/vision/dino_vits16/offline/bach.yaml

@@ -1,4 +1,5 @@
 ---
+n_runs: 5
 trainer:
   class_path: eva.Trainer
   init_args:
@@ -18,7 +19,7 @@ trainer:
       - class_path: pytorch_lightning.callbacks.EarlyStopping
         init_args:
           min_delta: 0
-          patience: 500
+          patience: 800
           monitor: *MONITOR_METRIC
           mode: *MONITOR_METRIC_MODE
       - class_path: eva.callbacks.EmbeddingsWriter
@@ -36,7 +37,7 @@ trainer:
                 model: dino_vits16
                 pretrained: ${oc.env:PRETRAINED, true}
     logger:
-      - class_path: pytorch_lightning.loggers.TensorBoardLogger
+      - class_path: eva.loggers.TensorBoardLogger
         init_args:
           save_dir: *LIGHTNING_ROOT
           name: ""

configs/vision/dino_vits16/offline/crc_he.yaml

@@ -1,4 +1,5 @@
 ---
+n_runs: 5
 trainer:
   class_path: eva.Trainer
   init_args:
@@ -18,7 +19,7 @@ trainer:
       - class_path: pytorch_lightning.callbacks.EarlyStopping
         init_args:
           min_delta: 0
-          patience: 100
+          patience: 48
           monitor: *MONITOR_METRIC
           mode: *MONITOR_METRIC_MODE
       - class_path: eva.callbacks.EmbeddingsWriter
@@ -36,7 +37,7 @@ trainer:
                 model: dino_vits16
                 pretrained: ${oc.env:PRETRAINED, true}
     logger:
-      - class_path: pytorch_lightning.loggers.TensorBoardLogger
+      - class_path: eva.loggers.TensorBoardLogger
         init_args:
           save_dir: *LIGHTNING_ROOT
           name: ""

configs/vision/dino_vits16/offline/crc_he_nonorm.yaml

@@ -1,4 +1,5 @@
 ---
+n_runs: 5
 trainer:
   class_path: eva.Trainer
   init_args:
@@ -18,7 +19,7 @@ trainer:
       - class_path: pytorch_lightning.callbacks.EarlyStopping
         init_args:
           min_delta: 0
-          patience: 100
+          patience: 48
           monitor: *MONITOR_METRIC
           mode: *MONITOR_METRIC_MODE
       - class_path: eva.callbacks.EmbeddingsWriter
@@ -36,7 +37,7 @@ trainer:
                 model: dino_vits16
                 pretrained: ${oc.env:PRETRAINED, true}
     logger:
-      - class_path: pytorch_lightning.loggers.TensorBoardLogger
+      - class_path: eva.loggers.TensorBoardLogger
         init_args:
           save_dir: *LIGHTNING_ROOT
           name: ""

configs/vision/dino_vits16/offline/patch_camelyon.yaml

@@ -1,4 +1,5 @@
 ---
+n_runs: 5
 trainer:
   class_path: eva.Trainer
   init_args:
@@ -18,7 +19,7 @@ trainer:
       - class_path: pytorch_lightning.callbacks.EarlyStopping
         init_args:
           min_delta: 0
-          patience: 100
+          patience: 25
           monitor: *MONITOR_METRIC
           mode: *MONITOR_METRIC_MODE
       - class_path: eva.callbacks.EmbeddingsWriter
@@ -37,7 +38,7 @@ trainer:
                 model: dino_vits16
                 pretrained: ${oc.env:PRETRAINED, true}
     logger:
-      - class_path: pytorch_lightning.loggers.TensorBoardLogger
+      - class_path: eva.loggers.TensorBoardLogger
         init_args:
           save_dir: *LIGHTNING_ROOT
           name: ""

configs/vision/dino_vits16/online/bach.yaml

@@ -22,7 +22,7 @@ trainer:
           monitor: *MONITOR_METRIC
           mode: *MONITOR_METRIC_MODE
     logger:
-      - class_path: pytorch_lightning.loggers.TensorBoardLogger
+      - class_path: eva.loggers.TensorBoardLogger
         init_args:
           save_dir: *LIGHTNING_ROOT
           name: ""

configs/vision/dino_vits16/online/patch_camelyon.yaml

@@ -22,7 +22,7 @@ trainer:
           monitor: *MONITOR_METRIC
           mode: *MONITOR_METRIC_MODE
     logger:
-      - class_path: pytorch_lightning.loggers.TensorBoardLogger
+      - class_path: eva.loggers.TensorBoardLogger
         init_args:
           save_dir: *LIGHTNING_ROOT
           name: ""

configs/vision/tests/offline/patch_camelyon.yaml

@@ -32,7 +32,7 @@ trainer:
           monitor: &MONITOR_METRIC val/BinaryAccuracy
           mode: &MONITOR_METRIC_MODE max
     logger:
-      - class_path: pytorch_lightning.loggers.TensorBoardLogger
+      - class_path: eva.loggers.TensorBoardLogger
         init_args:
           save_dir: *LIGHTNING_ROOT
           name: ""

configs/vision/tests/offline/patches.yaml

@@ -9,7 +9,7 @@ trainer:
         init_args:
           logging_interval: epoch
     logger:
-      - class_path: pytorch_lightning.loggers.TensorBoardLogger
+      - class_path: eva.loggers.TensorBoardLogger
         init_args:
           save_dir: *OUTPUT_DIR
           name: ""

configs/vision/tests/offline/slides.yaml

@@ -9,7 +9,7 @@ trainer:
         init_args:
           logging_interval: epoch
     logger:
-      - class_path: pytorch_lightning.loggers.TensorBoardLogger
+      - class_path: eva.loggers.TensorBoardLogger
         init_args:
           save_dir: *OUTPUT_DIR
           name: ""

configs/vision/tests/online/patch_camelyon.yaml

@@ -1,11 +1,19 @@
 ---
+n_runs: 2
 trainer:
   class_path: eva.Trainer
   init_args:
     default_root_dir: &OUTPUT_DIR ${oc.env:OUTPUT_DIR, logs/dino_vits16/patch_camelyon}
     max_epochs: &MAX_EPOCHS 1
     limit_train_batches: 2
     limit_val_batches: 2
+    logger:
+      - class_path: eva.loggers.TensorBoardLogger
+        init_args:
+          save_dir: *OUTPUT_DIR
+      - class_path: eva.loggers.CSVLogger
+        init_args:
+          save_dir: *OUTPUT_DIR
 model:
   class_path: eva.HeadModule
   init_args:

src/eva/callbacks/writers/embeddings.py

@@ -98,6 +98,8 @@ def on_predict_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule) ->
         self._write_queue.put(None)
         self._write_process.join()
         logger.info(f"Predictions and manifest saved to {self._output_dir}")
+        self._write_process = None  # type: ignore
+        self._write_queue = None  # type: ignore
 
     def _initialize_write_process(self) -> None:
         self._write_queue = multiprocessing.Queue()

src/eva/interface/interface.py

@@ -1,9 +1,18 @@
 """Main interface class."""
 
+import copy
+import os
+from datetime import datetime
+
+import pytorch_lightning as pl
+from loguru import logger
+from pytorch_lightning.callbacks import ModelCheckpoint
+
 from eva import trainers
 from eva.data import datamodules
 from eva.data.datamodules import schemas
 from eva.models import modules
+from eva.utils.recording import get_evaluation_id, record_results
 
 
 class Interface:
@@ -18,6 +27,7 @@ def fit(
         model: modules.ModelModule,
         data: datamodules.DataModule,
         trainer: trainers.Trainer,
+        n_runs: int = 1,
     ) -> None:
         """Perform model training and evaluation in place.
 
@@ -34,11 +44,24 @@ def fit(
             model: The model module.
             data: The data module.
             trainer: The trainer which processes the model and data.
+            n_runs: The number of runs to perform.
         """
-        trainer.fit(model=model, datamodule=data)
-        trainer.validate(datamodule=data)
-        if data.datasets.test is not None:
-            trainer.test(datamodule=data)
+        evaluation_id = get_evaluation_id()
+
+        for run_id in range(n_runs):
+            _trainer = copy.deepcopy(trainer)
+            _model = copy.deepcopy(model)
+            log_dir = os.path.join(_trainer.default_root_dir, evaluation_id, f"run_{run_id}")
+            _adapt_log_dirs(_trainer, log_dir)
+
+            start_time = datetime.now()
+            pl.seed_everything(run_id + 3, workers=True)
+
+            evaluation_results = _fit_validate_test(_trainer, _model, data)
+
+            end_time = datetime.now()
+            results_path = os.path.join(log_dir, "results.json")
+            record_results(evaluation_results, results_path, start_time, end_time)
 
     def predict(
         self,
@@ -80,3 +103,51 @@ def predict_fit(
         """
         self.predict(model=model, data=data, trainer=trainer)
         self.fit(model=model, data=data, trainer=trainer)
+
+
+def _fit_validate_test(
+    trainer: trainers.Trainer,
+    model: modules.ModelModule,
+    data: datamodules.DataModule,
+) -> dict:
+    """Combines the fit and validate commands in one method.
+
+    Helper method to perform the following three steps:
+    1. fit: training the model using the provided data.
+    2. validate: evaluating the model using the validation data.
+    3. test: evaluating the model using the test data. (if available)
+
+    Args:
+        model: The model module.
+        data: The data module.
+        trainer: The trainer which processes the model and data.
+    """
+    trainer.fit(model=model, datamodule=data)
+    evaluation_results = {"val": trainer.validate(datamodule=data)}
+    if data.datasets.test is not None:
+        evaluation_results["test"] = trainer.test(datamodule=data)
+    return evaluation_results
+
+
+def _adapt_log_dirs(trainer, log_dir: str) -> None:
+    """Sets the log directory for the logger, trainer and callbacks.
+
+    Args:
+        trainer: The trainer instance.
+        log_dir: The log directory.
+    """
+    for train_logger in trainer.loggers:
+        try:
+            train_logger.log_dir = log_dir
+        except Exception:
+            logger.warning(f"Could not set log_dir for logger {train_logger}")
+
+    trainer.log_dir = log_dir
+    if len(trainer.callbacks) > 0:
+        model_checkpoint_callbacks = [
+            c for c in trainer.callbacks if isinstance(c, ModelCheckpoint)
+        ]
+        if len(model_checkpoint_callbacks) > 0:
+            model_checkpoint_callbacks[0].dirpath = os.path.join(log_dir, "checkpoints")
+        else:
+            logger.warning("No ModelCheckpoint callback found in trainer.callbacks")

src/eva/loggers/__init__.py

@@ -0,0 +1,5 @@
+"""Loggers API."""
+
+from eva.loggers.lightning import CSVLogger, TensorBoardLogger
+
+__all__ = ["CSVLogger", "TensorBoardLogger"]

src/eva/loggers/lightning.py

@@ -0,0 +1,44 @@
+"""Custom logger classes for PyTorch Lightning."""
+
+from pytorch_lightning import loggers
+from typing_extensions import override
+
+
+class BaseLogger(loggers.Logger):
+    """Base logger class."""
+
+    def __init__(self, *args, **kwargs):
+        """Initializes the BaseLogger instance.
+
+        Overwrites the parent class to allow for custom log_dir setting.
+        """
+        super().__init__(*args, **kwargs)
+        self._log_dir = None
+
+    @property
+    @override
+    def log_dir(self) -> str:
+        if self._log_dir is not None:
+            return self._log_dir
+        else:
+            return super().log_dir  # type: ignore
+
+    @log_dir.setter
+    def log_dir(self, value):
+        self._log_dir = value
+
+
+class TensorBoardLogger(BaseLogger, loggers.TensorBoardLogger):
+    """TensorBoard logger class."""
+
+    def __init__(self, *args, **kwargs):
+        """Initializes the TensorBoardLogger instance."""
+        super().__init__(*args, **kwargs)
+
+
+class CSVLogger(BaseLogger, loggers.CSVLogger):
+    """CSV logger class."""
+
+    def __init__(self, *args, **kwargs):
+        """Initializes the CSVLogger instance."""
+        super().__init__(*args, **kwargs)

src/eva/trainers/trainer.py

@@ -1,6 +1,32 @@
 """Core trainer module."""
 
+from typing import Optional
+
 from pytorch_lightning import trainer
+from typing_extensions import override
+
+from eva.utils.recording import get_evaluation_id
+
+
+class Trainer(trainer.Trainer):
+    """Core Trainer class."""
+
+    def __init__(self, **kwargs):
+        """Initializes a new Trainer instance."""
+        super(Trainer, self).__init__(**kwargs)
+        self.evaluation_id = get_evaluation_id()
+        self._log_dir = None
+        self.i = 0
+
+    @property
+    @override
+    def log_dir(self) -> Optional[str]:
+        """Overrides the log_dir getter from parent class."""
+        if self._log_dir is not None:
+            return self._log_dir
+        else:
+            return super().log_dir
 
-Trainer = trainer.Trainer
-"""Core trainer class."""
+    @log_dir.setter
+    def log_dir(self, value):
+        self._log_dir = value