""" BEWARE MLP Training Pipeline Project: BEWARE Summer Research Internship Institution: University of Bath Description ----------- This script trains and evaluates a multilayer perceptron (MLP) for predicting the 2% exceedance wave run-up (R2) using the combined Scott et al. and BEWARE-2 datasets. The workflow implements: 1. Group K-Fold cross-validation using profile IDs to avoid information leakage. 2. Hyperparameter optimisation using Optuna with the Hyperband pruner. 3. Retraining of the best-performing architecture on the complete development dataset. 4. Prediction on the held-out test fold. 5. Saving of the trained model, scaler, predictions, optimisation history and performance metrics. Each fold is executed independently via python 1_nested_hyperband_cv.py --fold N allowing multiple folds to be run in parallel on an HPC system. """ # BEWARE MLP # Nested Group K-Fold Cross Validation # Optuna + Hyperband (ASHA) # # PART 1 # Imports # Configuration # Device # Reproducibility # Data loading # Dataset classes # MLP model # ----------- # IMPORTS # ----------- import os import gc import copy import json import random import pickle import argparse import numpy as np import pandas as pd import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader from sklearn.model_selection import ( GroupKFold, GroupShuffleSplit ) from sklearn.preprocessing import StandardScaler from sklearn.metrics import ( mean_squared_error, mean_absolute_error, r2_score ) import optuna from optuna.pruners import HyperbandPruner # ----------------------- # COMMAND LINE ARGUMENTS # ----------------------- parser = argparse.ArgumentParser( description="Run a single fold of nested GroupKFold." ) parser.add_argument( "--fold", type=int, required=True, help="Fold number to run (1-5)" ) args = parser.parse_args() REQUESTED_FOLD = args.fold # -------------- # CONFIGURATION # -------------- from pathlib import Path BASE_DIR = Path(__file__).resolve().parent CONFIG = { "data_file": BASE_DIR / "BEWARE_combined_randomised.parquet", "output_folder": BASE_DIR / "Nested_CV_Results", "target": "r2", "group_column": "profileIds", "random_seed": 42, # -------------------------------------------------------- # Cross Validation # -------------------------------------------------------- "outer_folds": 5, "validation_fraction": 0.10, # -------------------------------------------------------- # Hyperband # -------------------------------------------------------- "n_trials": 50, "max_epochs": 200, "patience": 30, # Development mode. # When enabled, the number of trials and training epochs are reduced # to allow rapid testing on a local workstation. "fast_mode": False } if not (1 <= REQUESTED_FOLD <= CONFIG["outer_folds"]): raise ValueError( f"Fold must be between 1 and {CONFIG['outer_folds']}" ) # ----------- # FAST MODE # ----------- if CONFIG["fast_mode"]: CONFIG["n_trials"] = 5 CONFIG["max_epochs"] = 30 CONFIG["patience"] = 5 # DON'T change outer_folds # ----------- # DEVICE # ----------- DEVICE = torch.device( "cuda" if torch.cuda.is_available() else "cpu" ) print("\n==================================================") print("BEWARE Nested Hyperband Cross Validation") print("==================================================") print("Device:", DEVICE) # ---------------- # REPRODUCIBILITY # ---------------- def set_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) set_seed(CONFIG["random_seed"]) # ----------------- # OUTPUT DIRECTORY # ----------------- os.makedirs( CONFIG["output_folder"], exist_ok=True ) # -------------- # LOAD DATASET # -------------- print("\nLoading dataset...") df = pd.read_parquet( CONFIG["data_file"] ) print(f"Dataset shape : {df.shape}") # ---------------- # FEATURE COLUMNS # ---------------- REMOVE_COLUMNS = [ CONFIG["target"], CONFIG["group_column"], "ID", "LoadingCondition", "dataset" ] FEATURE_COLUMNS = [ column for column in df.columns if column not in REMOVE_COLUMNS ] print(f"Input features : {len(FEATURE_COLUMNS)}") # ------------------ # SMALL ARRAYS ONLY # ------------------ # The dataset is retained as a pandas DataFrame. # Feature matrices are extracted one fold at a time to reduce peak memory usage y = df[CONFIG["target"]].to_numpy( dtype=np.float32 ) groups = df[CONFIG["group_column"]].to_numpy() print(f"Samples : {len(df)}") print(f"Groups : {len(np.unique(groups))}") # ---------- # DATASETS # ---------- class BewareDataset(Dataset): """ PyTorch Dataset wrapping the feature matrix and target values for a single training or validation split. The feature matrix is created only for the current fold, which avoids allocating one very large array for the entire dataset. """ def __init__(self, X, y): self.X = X self.y = y def __len__(self): return len(self.X) def __getitem__(self, index): return ( torch.from_numpy( self.X[index] ).float(), torch.tensor( self.y[index], dtype=torch.float32 ).view(1) ) class InferenceDataset(Dataset): def __init__(self, X): self.X = X def __len__(self): return len(self.X) def __getitem__(self, index): return torch.from_numpy( self.X[index] ).float() # -------------- # MLP MODEL # -------------- class MLP(nn.Module): def __init__( self, input_size, hidden_layers, neurons, dropout, activation ): super().__init__() activation_layers = { "ReLU": nn.ReLU, "LeakyReLU": nn.LeakyReLU, "GELU": nn.GELU } layers = [] previous = input_size for _ in range(hidden_layers): layers.append( nn.Linear( previous, neurons ) ) layers.append( activation_layers[activation]() ) layers.append( nn.Dropout( dropout ) ) previous = neurons layers.append( nn.Linear( previous, 1 ) ) self.network = nn.Sequential( *layers ) def forward(self, x): return self.network(x) # -------------- # HYPERPARAMETER SEARCH SPACE # -------------- SEARCH_SPACE = { "hidden_layers": [2, 3, 4], "neurons": [64, 128, 256, 512], "batch_size": [64, 128, 256], "activations": [ "ReLU", "LeakyReLU", "GELU" ] } print("\nConfiguration") print("--------------------------------------------------") for key, value in CONFIG.items(): print(f"{key:25s}: {value}") print("--------------------------------------------------") # ------------------ # TRAINING FUNCTION # ------------------ def train_model( model, train_loader, validation_loader, optimiser, max_epochs, patience, trial=None ): criterion = nn.MSELoss() history = { "epoch": [], "train_loss": [], "validation_loss": [] } best_loss = np.inf best_state = None patience_counter = 0 for epoch in range(max_epochs): # ---------------------------------------------------- # Training # ---------------------------------------------------- model.train() train_losses = [] for xb, yb in train_loader: xb = xb.to(DEVICE) yb = yb.to(DEVICE) optimiser.zero_grad() prediction = model(xb) loss = criterion( prediction, yb ) loss.backward() optimiser.step() train_losses.append(loss.item()) train_loss = np.mean(train_losses) # ---------------------------------------------------- # Validation # ---------------------------------------------------- model.eval() validation_losses = [] with torch.no_grad(): for xb, yb in validation_loader: xb = xb.to(DEVICE) yb = yb.to(DEVICE) prediction = model(xb) validation_losses.append( criterion( prediction, yb ).item() ) validation_loss = np.mean(validation_losses) history["epoch"].append(epoch + 1) history["train_loss"].append(train_loss) history["validation_loss"].append(validation_loss) if trial is not None: trial.report( validation_loss, epoch ) if trial.should_prune(): raise optuna.TrialPruned() if validation_loss < best_loss: best_loss = validation_loss best_state = copy.deepcopy( model.state_dict() ) patience_counter = 0 else: patience_counter += 1 if patience_counter >= patience: break if best_state is not None: model.load_state_dict(best_state) return best_loss, history # ------------------- # OUTER GROUP K-FOLD # ------------------- outer_cv = GroupKFold( n_splits=CONFIG["outer_folds"] ) fold_results = [] for fold_number, (development_index, test_index) in enumerate( outer_cv.split( df, y, groups ), start=1 ): if fold_number != REQUESTED_FOLD: continue # Skip all folds except the one requested via --fold. print("\n=================================================") print(f"FOLD {fold_number}/{CONFIG['outer_folds']}") print("=================================================") fold_folder = os.path.join( CONFIG["output_folder"], f"fold_{fold_number:02d}" ) os.makedirs( fold_folder, exist_ok=True ) # -------------------------------------------------------- # Split development into train/validation # -------------------------------------------------------- development_groups = groups[development_index] splitter = GroupShuffleSplit( n_splits=1, test_size=CONFIG["validation_fraction"], random_state=CONFIG["random_seed"] ) train_local, validation_local = next( splitter.split( development_index, y[development_index], development_groups ) ) train_index = development_index[train_local] validation_index = development_index[validation_local] # -------------------------------------------------------- # Extract feature matrices for the current fold only. # This minimises memory usage during cross-validation. # -------------------------------------------------------- X_train = df.loc[ train_index, FEATURE_COLUMNS ].to_numpy(dtype=np.float32) y_train = y[train_index] X_validation = df.loc[ validation_index, FEATURE_COLUMNS ].to_numpy(dtype=np.float32) y_validation = y[validation_index] X_test = df.loc[ test_index, FEATURE_COLUMNS ].to_numpy(dtype=np.float32) y_test = y[test_index] # -------------------------------------------------------- # Scale # -------------------------------------------------------- scaler = StandardScaler() X_train = scaler.fit_transform( X_train ).astype(np.float32) X_validation = scaler.transform( X_validation ).astype(np.float32) X_test = scaler.transform( X_test ).astype(np.float32) # -------------------------------------------------------- # Optuna objective # -------------------------------------------------------- def objective(trial): hidden_layers = trial.suggest_categorical( "hidden_layers", SEARCH_SPACE["hidden_layers"] ) neurons = trial.suggest_categorical( "neurons", SEARCH_SPACE["neurons"] ) batch_size = trial.suggest_categorical( "batch_size", SEARCH_SPACE["batch_size"] ) activation = trial.suggest_categorical( "activation", SEARCH_SPACE["activations"] ) dropout = trial.suggest_float( "dropout", 0.0, 0.5 ) learning_rate = trial.suggest_float( "learning_rate", 1e-4, 1e-2, log=True ) weight_decay = trial.suggest_float( "weight_decay", 1e-6, 1e-3, log=True ) train_loader = DataLoader( BewareDataset( X_train, y_train ), batch_size=batch_size, shuffle=True ) validation_loader = DataLoader( BewareDataset( X_validation, y_validation ), batch_size=batch_size, shuffle=False ) model = MLP( input_size=X_train.shape[1], hidden_layers=hidden_layers, neurons=neurons, dropout=dropout, activation=activation ).to(DEVICE) optimiser = torch.optim.Adam( model.parameters(), lr=learning_rate, weight_decay=weight_decay ) validation_loss, history = train_model( model, train_loader, validation_loader, optimiser, CONFIG["max_epochs"], CONFIG["patience"], trial ) pd.DataFrame(history).to_csv( os.path.join( fold_folder, f"trial_{trial.number:03d}_history.csv" ), index=False ) return validation_loss # -------------------------------------------------------- # Hyperband # Hyperparameter optimisation using Optuna. # Hyperband prunes poorly performing trials early to reduce # computational cost while exploring the search space efficiently. # -------------------------------------------------------- print("Running Hyperband search...") study = optuna.create_study( study_name=f"fold_{fold_number}", direction="minimize", pruner=HyperbandPruner( min_resource=10, max_resource=CONFIG["max_epochs"], reduction_factor=3 ), storage=f"sqlite:///{CONFIG['output_folder']}/fold_{fold_number}.db", load_if_exists=False ) study.optimize( objective, n_trials=CONFIG["n_trials"] ) best_parameters = study.best_trial.params print( "Best validation loss:", study.best_trial.value ) # ------------------- # RETRAIN BEST MODEL # Retrain the best hyperparameter configuration using the # combined training and validation data before final testing. # ------------------- print("Retraining best model...") batch_size = best_parameters["batch_size"] train_validation_index = np.concatenate( [train_index, validation_index] ) del X_train del X_validation del y_train del y_validation gc.collect() X_train_full = df.loc[ train_validation_index, FEATURE_COLUMNS ].to_numpy(dtype=np.float32) y_train_full = y[train_validation_index] scaler = StandardScaler() X_train_full = scaler.fit_transform( X_train_full ).astype(np.float32) X_test = df.loc[ test_index, FEATURE_COLUMNS ].to_numpy(dtype=np.float32) X_test = scaler.transform( X_test ).astype(np.float32) train_loader = DataLoader( BewareDataset( X_train_full, y_train_full ), batch_size=batch_size, shuffle=True ) test_loader = DataLoader( InferenceDataset( X_test ), batch_size=batch_size, shuffle=False ) model = MLP( input_size=X_train_full.shape[1], hidden_layers=best_parameters["hidden_layers"], neurons=best_parameters["neurons"], dropout=best_parameters["dropout"], activation=best_parameters["activation"] ).to(DEVICE) optimiser = torch.optim.Adam( model.parameters(), lr=best_parameters["learning_rate"], weight_decay=best_parameters["weight_decay"] ) train_model( model, train_loader, train_loader, optimiser, CONFIG["max_epochs"], CONFIG["patience"] ) # -------------- # SAVE MODEL # -------------- torch.save( model.state_dict(), os.path.join( fold_folder, "best_model.pt" ) ) # -------------- # PREDICTION # -------------- print("Predicting...") model.eval() predictions = [] with torch.no_grad(): for xb in test_loader: xb = xb.to(DEVICE) prediction = model(xb) predictions.extend( prediction.squeeze(1) .cpu() .numpy() ) predictions = np.asarray(predictions) # -------------- # METRICS # -------------- rmse = np.sqrt( mean_squared_error( y_test, predictions ) ) mae = mean_absolute_error( y_test, predictions ) r2 = r2_score( y_test, predictions ) print(f"RMSE : {rmse:.5f}") print(f"MAE : {mae:.5f}") print(f"R² : {r2:.5f}") fold_results.append({ "fold": fold_number, "rmse": rmse, "mae": mae, "r2": r2 }) # ----------------- # SAVE PREDICTIONS # -------------- # Information for the test samples # Save predictions together with sample identifiers to enable # subsequent analysis against environmental parameters. test_ids = df.loc[test_index, "ID"].to_numpy() test_profiles = df.loc[test_index, "profileIds"].to_numpy() prediction_df = pd.DataFrame({ "ID": test_ids, "profileIds": test_profiles, "Observed": y_test, "Predicted": predictions }) prediction_df.to_csv( os.path.join( fold_folder, "predictions.csv" ), index=False ) # --------------------- # SAVE BEST PARAMETERS # --------------------- with open( os.path.join( fold_folder, "best_parameters.json" ), "w" ) as file: json.dump( best_parameters, file, indent=4 ) # ------------------ # SAVE ALL TRIALS # ----------------- study.trials_dataframe().to_csv( os.path.join( fold_folder, "all_trials.csv" ), index=False ) # -------------- # SAVE SCALER # -------------- with open( os.path.join( fold_folder, "scaler.pkl" ), "wb" ) as file: pickle.dump( scaler, file ) # ---------------- # CLEAN MEMORY # ---------------- del model del optimiser del train_loader del test_loader del X_train_full del X_test del y_train_full del y_test del scaler gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() break # ----------------- # OVERALL RESULTS # ----------------- results_df = pd.DataFrame( fold_results ) results_df.to_csv( CONFIG["output_folder"] / f"results_fold{REQUESTED_FOLD}.csv", index=False ) summary = { "RMSE Mean": float( results_df["rmse"].mean() ), "RMSE SD": float( results_df["rmse"].std() ), "MAE Mean": float( results_df["mae"].mean() ), "MAE SD": float( results_df["mae"].std() ), "R2 Mean": float( results_df["r2"].mean() ), "R2 SD": float( results_df["r2"].std() ) } with open( os.path.join( CONFIG["output_folder"], f"summary_fold{REQUESTED_FOLD}.json" ), "w" ) as file: json.dump( summary, file, indent=4 ) print("\n==================================================") print("Cross-validation complete") print("==================================================") print(results_df) print("\nSummary") for key, value in summary.items(): print(f"{key:15s}: {value:.5f}")