Merge branch 'main' into architecture_refactoring

Author: robin-janssen

codes/benchmark/__init__.py

@@ -73,6 +73,7 @@
     read_yaml_config,
     save_table_csv,
     write_metrics_to_yaml,
+    measure_inference_time,
 )
 
 __all__ = [
@@ -146,4 +147,5 @@
     "make_comparison_csv",
     "save_table_csv",
     "get_model_config",
+    "measure_inference_time",
 ]

codes/benchmark/bench_fcts.py

@@ -1,4 +1,3 @@
-import time
 from contextlib import redirect_stdout
 from typing import Any
 
@@ -40,6 +39,7 @@
     get_model_config,
     get_surrogate,
     make_comparison_csv,
+    measure_inference_time,
     measure_memory_footprint,
     save_table_csv,
     write_metrics_to_yaml,
@@ -349,13 +349,12 @@ def time_inference(
     n_runs: int = 5,
 ) -> dict[str, Any]:
     """
-    Time the inference of the surrogate model.
+    Time the inference of the surrogate model (full version with metrics).
 
     Args:
         model: Instance of the surrogate model class.
         surr_name (str): The name of the surrogate model.
         test_loader (DataLoader): The DataLoader object containing the test data.
-        timesteps (np.ndarray): The timesteps array.
         conf (dict): The configuration dictionary.
         n_test_samples (int): The number of test samples.
         n_runs (int, optional): Number of times to run the inference for timing.
@@ -366,35 +365,19 @@ def time_inference(
     training_id = conf["training_id"]
     model.load(training_id, surr_name, model_identifier=f"{surr_name.lower()}_main")
 
-    # Run inference multiple times and record the durations
-    inference_times = []
-    for _ in range(n_runs):
-        # _, _ = model.predict(data_loader=test_loader)
-        total_time = 0
-        with torch.inference_mode():
-            for inputs in test_loader:
-                start_time = time.perf_counter()
-                _, _ = model.forward(inputs)
-                end_time = time.perf_counter()
-                total_time += end_time - start_time
-        # total_time /= n_test_samples
-        inference_times.append(total_time)
-
-    # Calculate metrics
+    inference_times = measure_inference_time(model, test_loader, n_runs=n_runs)
+
     mean_inference_time = np.mean(inference_times)
     std_inference_time = np.std(inference_times)
 
-    # Store metrics
-    timing_metrics = {
+    return {
         "mean_inference_time_per_run": mean_inference_time,
         "std_inference_time_per_run": std_inference_time,
         "num_predictions": n_test_samples,
         "mean_inference_time_per_prediction": mean_inference_time / n_test_samples,
         "std_inference_time_per_prediction": std_inference_time / n_test_samples,
     }
 
-    return timing_metrics
-
 
 def evaluate_compute(
     model, surr_name: str, test_loader: DataLoader, conf: dict

codes/benchmark/bench_utils.py

@@ -8,10 +8,13 @@
 import numpy as np
 import torch
 import yaml
+from torch.utils.data import DataLoader
 
 from codes.surrogates import SurrogateModel, surrogate_classes
 from codes.utils import read_yaml_config
 
+import time
+
 
 def check_surrogate(surrogate: str, conf: dict) -> None:
     """
@@ -699,3 +702,32 @@ def get_model_config(surr_name: str, config: dict) -> dict:
         model_config = {}
 
     return model_config
+
+
+def measure_inference_time(
+    model,
+    test_loader: DataLoader,
+    n_runs: int = 5,
+) -> list[float]:
+    """
+    Measure total inference time over a DataLoader across multiple runs.
+
+    Args:
+        model: Model instance with a `.forward()` method.
+        test_loader (DataLoader): Loader with test data.
+        n_runs (int): Number of repeated runs for averaging.
+
+    Returns:
+        list[float]: List of total inference times per run (in seconds).
+    """
+    inference_times = []
+    for _ in range(n_runs):
+        total_time = 0
+        with torch.inference_mode():
+            for inputs in test_loader:
+                start_time = time.perf_counter()
+                _, _ = model.forward(inputs)
+                end_time = time.perf_counter()
+                total_time += end_time - start_time
+        inference_times.append(total_time)
+    return inference_times

codes/surrogates/DeepONet/deeponet.py

@@ -286,6 +286,7 @@ def fit(
         epochs: int,
         position: int = 0,
         description: str = "Training DeepONet",
+        multi_objective: bool = False,
     ) -> None:
         """
         Train the MultiONet model.
@@ -296,6 +297,8 @@ def fit(
             epochs (int, optional): The number of epochs to train the model.
             position (int): The position of the progress bar.
             description (str): The description for the progress bar.
+            multi_objective (bool): Whether multi-objective optimization is used. 
+                                    If True, trial.report is not used (not supported by Optuna).
 
         Returns:
             None. The training loss, test loss, and MAE are stored in the model.
@@ -335,7 +338,8 @@ def fit(
                 progress_bar.set_postfix(postfix)
 
                 # Report the loss to Optuna and check for pruning
-                if self.optuna_trial is not None:
+                if self.optuna_trial is not None and not multi_objective:
+
                     self.optuna_trial.report(test_losses[index], epoch)
                     if self.optuna_trial.should_prune():
                         raise optuna.TrialPruned()

codes/surrogates/FCNN/fcnn.py

@@ -159,7 +159,23 @@ def fit(
         epochs: int,
         position: int = 0,
         description: str = "Training FullyConnected",
+        multi_objective: bool = False,
     ) -> None:
+        """
+        Train the FullyConnected model.
+
+        Args:
+            train_loader (DataLoader): The DataLoader object containing the training data.
+            test_loader (DataLoader): The DataLoader object containing the test data.
+            epochs (int, optional): The number of epochs to train the model.
+            position (int): The position of the progress bar.
+            description (str): The description for the progress bar.
+            multi_objective (bool): Whether multi-objective optimization is used.
+                                    If True, trial.report is not used (not supported by Optuna).
+
+        Returns:
+            None. The training loss, test loss, and MAE are stored in the model.
+        """
         self.n_train_samples = int(len(train_loader.dataset) / self.n_timesteps)
         # criterion = nn.MSELoss(reduction="sum")
         criterion = nn.MSELoss()
@@ -196,7 +212,7 @@ def fit(
                 progress_bar.set_postfix(postfix)
 
                 # Report the test loss to Optuna
-                if self.optuna_trial is not None:
+                if self.optuna_trial is not None and not multi_objective:
                     self.optuna_trial.report(test_losses[index], step=epoch)
                     if self.optuna_trial.should_prune():
                         raise optuna.TrialPruned()

codes/surrogates/LatentNeuralODE/latent_neural_ode.py

@@ -145,6 +145,7 @@ def fit(
         epochs: int,
         position: int = 0,
         description: str = "Training LatentNeuralODE",
+        multi_objective: bool = False,
     ) -> None:
         """
         Fits the model to the training data. Sets the train_loss and test_loss attributes.
@@ -156,6 +157,9 @@ def fit(
             epochs (int | None): The number of epochs to train the model. If None, uses the value from the config.
             position (int): The position of the progress bar.
             description (str): The description for the progress bar.
+            multi_objective (bool): Whether multi-objective optimization is used.
+                                    If True, trial.report is not used (not supported by Optuna).
+
         """
         # optimizer = Adam(self.model.parameters(), lr=self.config.learning_rate)
         optimizer = AdamWScheduleFree(
@@ -211,7 +215,7 @@ def fit(
                     progress_bar.set_postfix(postfix)
 
                     # Report loss to Optuna and prune if necessary
-                    if self.optuna_trial is not None:
+                    if self.optuna_trial is not None and not multi_objective:
                         self.optuna_trial.report(test_losses[index], step=epoch)
                         if self.optuna_trial.should_prune():
                             raise optuna.TrialPruned()

codes/surrogates/LatentPolynomial/latent_poly.py

@@ -144,16 +144,19 @@ def fit(
         epochs: int,
         position: int = 0,
         description: str = "Training LatentPoly",
+        multi_objective: bool = False,
     ) -> None:
         """
         Fit the model to the training data.
 
         Args:
-            train_loader (DataLoader): Training data loader.
-            test_loader (DataLoader): Test data loader.
-            epochs (int): Number of training epochs.
-            position (int): Progress bar position.
-            description (str): Description for the progress bar.
+            train_loader (DataLoader): The data loader for the training data.
+            test_loader (DataLoader): The data loader for the test data.
+            epochs (int | None): The number of epochs to train the model. If None, uses the value from the config.
+            position (int): The position of the progress bar.
+            description (str): The description for the progress bar.
+            multi_objective (bool): Whether multi-objective optimization is used.
+                                    If True, trial.report is not used (not supported by Optuna).
         """
         optimizer = AdamWScheduleFree(
             self.model.parameters(), lr=self.config.learning_rate
@@ -199,7 +202,8 @@ def fit(
                         }
                     )
 
-                    if self.optuna_trial is not None:
+                    # Report loss to Optuna and prune if necessary
+                    if self.optuna_trial is not None and not multi_objective:
                         self.optuna_trial.report(test_losses[index], step=epoch)
                         if self.optuna_trial.should_prune():
                             raise optuna.TrialPruned()

codes/tune/optuna_fcts.py

@@ -2,12 +2,17 @@
 import queue
 from distutils.util import strtobool
 
+import numpy as np
 import optuna
 import torch
 import torch.nn as nn
 import yaml
 
-from codes.benchmark.bench_utils import get_model_config, get_surrogate
+from codes.benchmark.bench_utils import (
+    get_model_config,
+    get_surrogate,
+    measure_inference_time,
+)
 from codes.utils import check_and_load_data, make_description, set_random_seeds
 from codes.utils.data_utils import get_data_subset
 
@@ -132,9 +137,10 @@ def objective(trial):
 
 def training_run(
     trial: optuna.Trial, device: str, config: dict, study_name: str
-) -> float:
+) -> float | tuple[float, float]:
     """
     Run the training for a single Optuna trial and return the loss.
+    In multi-objective mode, also returns the mean inference time.
 
     Args:
         trial (optuna.Trial): Optuna trial object.
@@ -143,9 +149,11 @@ def training_run(
         study_name (str): Name of the study.
 
     Returns:
-        float: Loss value.
+        float: Loss value in single objective mode.
+        tuple[float, float]: (loss, mean_inference_time) in multi objective mode.
     """
 
+    download_data(config["dataset"]["name"], verbose=False)
     train_data, test_data, val_data, timesteps, _, data_params, _ = check_and_load_data(
         config["dataset"]["name"],
         verbose=False,
@@ -195,6 +203,7 @@ def training_run(
         epochs=config["epochs"],
         position=pos,
         description=description,
+        multi_objective=config["multi_objective"],
     )
 
     criterion = torch.nn.MSELoss()
@@ -210,4 +219,11 @@ def training_run(
         base_dir="",
         training_id=savepath,
     )
-    return loss
+
+    # Check if we're running multi-objective optimisation
+    if config["multi_objective"]:
+        # Measure inference time
+        inference_times = measure_inference_time(model, test_loader)
+        return loss, np.mean(inference_times)
+    else:
+        return loss

codes/utils/data_utils.py

@@ -483,20 +483,22 @@ def update_to(self, b=1, bsize=1, tsize=None):
         self.update(b * bsize - self.n)
 
 
-def download_data(dataset_name: str, path: str | None = None):
+def download_data(dataset_name: str, path: str | None = None, verbose: bool = True):
     """
     Download the specified dataset if it is not present, with a progress bar.
     Args:
         dataset_name (str): The name of the dataset.
         path (str, optional): The path to save the dataset. If None, the default data directory is used.
+        verbose (bool): Whether to print information about the download progress.
     """
     data_path = (
         os.path.abspath(f"datasets/{dataset_name.lower()}/data.hdf5")
         if path is None
         else os.path.abspath(path)
     )
     if os.path.isfile(data_path):
-        print(f"Dataset '{dataset_name}' already downloaded at {data_path}.")
+        if verbose:
+            print(f"Dataset '{dataset_name}' already exists at {data_path}.")
         return
 
     with open("datasets/data_sources.yaml", "r", encoding="utf-8") as file: