442 lines
21 KiB
Python
442 lines
21 KiB
Python
import argparse
|
|
import logging
|
|
import sys
|
|
import warnings # Import the warnings module
|
|
|
|
import copy # For deep copying config
|
|
from pathlib import Path
|
|
import time
|
|
|
|
import pandas as pd
|
|
import torch
|
|
import yaml # Added for saving best config
|
|
|
|
import optuna
|
|
import pytorch_lightning as pl
|
|
from pytorch_lightning.callbacks import EarlyStopping
|
|
|
|
# Import necessary components from the forecasting_model package
|
|
from forecasting_model.utils.forecast_config_model import MainConfig
|
|
from forecasting_model.data_processing import (
|
|
prepare_fold_data_and_loaders,
|
|
split_data_classic
|
|
)
|
|
from forecasting_model.train.model import LSTMForecastLightningModule
|
|
from forecasting_model.train.classic import run_classic_training
|
|
|
|
|
|
# Import helper functions from forecasting_model_run.py
|
|
from forecasting_model.utils.helper import load_config, set_seeds
|
|
|
|
# Import the data processing functions
|
|
from forecasting_model.data_processing import load_raw_data
|
|
|
|
# --- Suppress specific PL warnings about logger=True with no logger ---
|
|
# This is expected behavior in optuna_run.py where logger=False is intentional
|
|
warnings.filterwarnings(
|
|
"ignore",
|
|
message=".*You called `self.log.*logger=True.*no logger configured.*",
|
|
category=UserWarning,
|
|
module="pytorch_lightning.core.module"
|
|
)
|
|
|
|
# Silence overly verbose libraries if needed
|
|
mpl_logger = logging.getLogger('matplotlib')
|
|
mpl_logger.setLevel(logging.WARNING)
|
|
pil_logger = logging.getLogger('PIL')
|
|
pil_logger.setLevel(logging.WARNING)
|
|
pl_logger = logging.getLogger('pytorch_lightning')
|
|
pl_logger.setLevel(logging.WARNING)
|
|
|
|
# --- Basic Logging Setup ---
|
|
logging.basicConfig(level=logging.INFO,
|
|
format='%(asctime)s - %(levelname)-7s - %(message)s',
|
|
datefmt='%H:%M:%S')
|
|
# Get the root logger
|
|
logger = logging.getLogger()
|
|
|
|
|
|
# --- Argument Parsing (Simplified) ---
|
|
def parse_arguments():
|
|
"""Parses command-line arguments for Optuna HPO."""
|
|
parser = argparse.ArgumentParser(
|
|
description="Run Hyperparameter Optimization using Optuna.",
|
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter
|
|
)
|
|
parser.add_argument(
|
|
'-c', '--config', type=str, default='forecasting_config.yaml',
|
|
help="Path to the YAML configuration file containing HPO settings."
|
|
)
|
|
parser.add_argument(
|
|
'--output-dir', type=str, default=None,
|
|
help="Override output directory specified in the configuration file."
|
|
)
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
# --- Optuna Objective Function ---
|
|
def objective(
|
|
trial: optuna.Trial,
|
|
base_config: MainConfig,
|
|
df: pd.DataFrame,
|
|
) -> float: # Ensure it returns a float
|
|
"""
|
|
Optuna single-objective function using a classic train/val/test split.
|
|
|
|
Returns:
|
|
- Validation score from the classic split (minimize).
|
|
"""
|
|
logger.info(f"\n--- Starting Optuna Trial {trial.number} ---")
|
|
trial_start_time = time.perf_counter()
|
|
|
|
# Get HPO settings
|
|
hpo_config = base_config.optuna
|
|
# Use the specific metric name from config for validation monitoring
|
|
validation_metric_monitor = hpo_config.metric_to_optimize
|
|
# The objective will be based on the validation metric
|
|
# Hardcode minimization for the single objective
|
|
monitor_mode = "min"
|
|
worst_value = float('inf') # Represents a poor validation score
|
|
|
|
# --- 1. Suggest Hyperparameters ---
|
|
try:
|
|
# Deep copy the base config dictionary to modify for the trial
|
|
trial_config_dict = copy.deepcopy(base_config.model_dump(mode='python')) # Use mode='python' for easier modification
|
|
except Exception as e:
|
|
logger.error(f"Trial {trial.number}: Failed to deep copy base configuration: {e}", exc_info=True)
|
|
# Return worst value
|
|
return worst_value
|
|
|
|
# ----- Suggest Hyperparameters -----
|
|
# (Suggest parameters as before, modifying trial_config_dict)
|
|
trial_config_dict['training']['learning_rate'] = trial.suggest_float('learning_rate', 1e-5, 1e-2, log=True)
|
|
trial_config_dict['training']['batch_size'] = trial.suggest_categorical('batch_size', [32, 64, 128, 256])
|
|
trial_config_dict['training']['loss_function'] = trial.suggest_categorical('loss_function', ['MSE', 'MAE'])
|
|
trial_config_dict['model']['hidden_size'] = trial.suggest_int('hidden_size', 18, 498, step=32)
|
|
trial_config_dict['model']['num_layers'] = trial.suggest_int('num_layers', 1, 8)
|
|
trial_config_dict['model']['dropout'] = trial.suggest_float('dropout', 0.0, 0.25, step=0.05)
|
|
trial_config_dict['features']['sequence_length'] = trial.suggest_int('sequence_length', 24, 168, step=12)
|
|
# Note: forecast_horizon is NOT tuned here, taken from base_config
|
|
trial_config_dict['features']['scaling_method'] = trial.suggest_categorical('scaling_method', ['standard', 'minmax', None])
|
|
use_configured_lags = trial.suggest_categorical('use_lags', [True, False])
|
|
if not use_configured_lags: trial_config_dict['features']['lags'] = []
|
|
use_configured_rolling = trial.suggest_categorical('use_rolling_windows', [True, False])
|
|
if not use_configured_rolling: trial_config_dict['features']['rolling_window_sizes'] = []
|
|
trial_config_dict['features']['use_time_features'] = trial.suggest_categorical('use_time_features', [True, False])
|
|
trial_config_dict['features']['sinus_curve'] = trial.suggest_categorical('sinus_curve', [True, False])
|
|
trial_config_dict['features']['cosin_curve'] = trial.suggest_categorical('cosin_curve', [True, False])
|
|
trial_config_dict['features']['fill_nan'] = trial.suggest_categorical('fill_nan', ['ffill', 'bfill', 0])
|
|
# ----- End of Hyperparameter Suggestions -----
|
|
|
|
# --- 2. Re-validate Trial Config ---
|
|
try:
|
|
# Transfer forecast horizon from base config (not tuned)
|
|
trial_config_dict['features']['forecast_horizon'] = base_config.features.forecast_horizon
|
|
trial_config = MainConfig(**trial_config_dict)
|
|
logger.info(f"Trial {trial.number} Parameters: {trial.params}")
|
|
except Exception as e:
|
|
logger.error(f"Trial {trial.number}: Invalid configuration generated: {e}")
|
|
return worst_value
|
|
|
|
# --- Early check for invalid sequence length / forecast horizon combination ---
|
|
try:
|
|
# Make sure forecast_horizon is a list before checking max()
|
|
if not isinstance(trial_config.features.forecast_horizon, list) or not trial_config.features.forecast_horizon:
|
|
raise ValueError("Trial config has invalid forecast_horizon list.")
|
|
min_data_for_sequence = trial_config.features.sequence_length + max(trial_config.features.forecast_horizon)
|
|
if min_data_for_sequence > len(df):
|
|
logger.warning(f"Trial {trial.number}: Skipped. sequence_length ({trial_config.features.sequence_length}) + "
|
|
f"max_horizon ({max(trial_config.features.forecast_horizon)}) "
|
|
f"exceeds data length ({len(df)}).")
|
|
# Optuna doesn't directly support skipping, so return worst values
|
|
return worst_value
|
|
except Exception as e:
|
|
logger.error(f"Trial {trial.number}: Error during pre-check: {e}", exc_info=True)
|
|
return worst_value
|
|
|
|
# --- 3. Run Classic Train/Test ---
|
|
logger.info(f"Trial {trial.number}: Starting Classic Run...")
|
|
validation_metric_value = worst_value # Initialize to worst
|
|
try:
|
|
n_samples = len(df)
|
|
val_frac = trial_config.cross_validation.val_size_fraction
|
|
test_frac = trial_config.cross_validation.test_size_fraction
|
|
train_idx_cl, val_idx_cl, test_idx_cl = split_data_classic(n_samples, val_frac, test_frac)
|
|
|
|
|
|
# Prepare data for classic split
|
|
train_loader_cl, val_loader_cl, test_loader_cl, target_scaler_cl, input_size_cl = prepare_fold_data_and_loaders(
|
|
full_df=df, train_idx=train_idx_cl, val_idx=val_idx_cl, test_idx=test_idx_cl,
|
|
target_col=trial_config.data.target_col, feature_config=trial_config.features,
|
|
train_config=trial_config.training, eval_config=trial_config.evaluation
|
|
)
|
|
|
|
# Initialize Model
|
|
model_cl = LSTMForecastLightningModule(
|
|
model_config=trial_config.model, train_config=trial_config.training,
|
|
input_size=input_size_cl, target_scaler=target_scaler_cl
|
|
)
|
|
|
|
# Callbacks (EarlyStopping and Pruning)
|
|
callbacks_cl = []
|
|
if trial_config.training.early_stopping_patience and trial_config.training.early_stopping_patience > 0:
|
|
callbacks_cl.append(EarlyStopping(
|
|
monitor=validation_metric_monitor, # Monitor validation metric
|
|
patience=trial_config.training.early_stopping_patience,
|
|
mode=monitor_mode, verbose=False
|
|
))
|
|
|
|
# Trainer for classic run
|
|
trainer_cl = pl.Trainer(
|
|
accelerator='gpu' if torch.cuda.is_available() else 'cpu', devices=1 if torch.cuda.is_available() else None,
|
|
max_epochs=trial_config.training.epochs, callbacks=callbacks_cl, logger=False, # logger=False as per original
|
|
enable_checkpointing=False, enable_progress_bar=False, enable_model_summary=False,
|
|
gradient_clip_val=getattr(trial_config.training, 'gradient_clip_val', None),
|
|
precision=getattr(trial_config.training, 'precision', 32),
|
|
)
|
|
|
|
# --- Train Model ---
|
|
logger.info(f"Trial {trial.number}: Fitting model on classic train/val split...")
|
|
trainer_cl.fit(model_cl, train_dataloaders=train_loader_cl, val_dataloaders=val_loader_cl)
|
|
|
|
# --- Get Best Validation Score ---
|
|
# Check early stopping callback first if it exists
|
|
best_score_tensor = None
|
|
if callbacks_cl and isinstance(callbacks_cl[0], EarlyStopping):
|
|
if hasattr(callbacks_cl[0], 'best_score') and callbacks_cl[0].best_score is not None:
|
|
best_score_tensor = callbacks_cl[0].best_score
|
|
elif callbacks_cl[0].stopped_epoch > 0 : # Early stopping triggered
|
|
logger.debug(f"Trial {trial.number}: Early stopping triggered, attempting to use last callback metric.")
|
|
|
|
# If early stopping didn't capture best score, use last metrics from trainer
|
|
if best_score_tensor is None:
|
|
metric_val = trainer_cl.callback_metrics.get(validation_metric_monitor)
|
|
if metric_val is not None:
|
|
best_score_tensor = metric_val # Use the last logged value
|
|
|
|
if best_score_tensor is None:
|
|
logger.warning(f"Trial {trial.number}: Metric '{validation_metric_monitor}' not found in callbacks or metrics. Using {worst_value}.")
|
|
validation_metric_value = worst_value
|
|
else:
|
|
validation_metric_value = best_score_tensor.item()
|
|
logger.info(f"Trial {trial.number}: Best val score ({validation_metric_monitor}) = {validation_metric_value:.4f}")
|
|
|
|
# Report intermediate value for pruning (if enabled)
|
|
trial.report(validation_metric_value, trainer_cl.current_epoch)
|
|
if trial.should_prune():
|
|
logger.info(f"Trial {trial.number}: Pruned.")
|
|
raise optuna.TrialPruned()
|
|
|
|
# Note: We don't run prediction/evaluation on the test set here,
|
|
# as the objective is based on validation performance.
|
|
# The test set evaluation will be done later for the best trial.
|
|
|
|
logger.info(f"Trial {trial.number}: Finished Classic Run in {time.perf_counter() - trial_start_time:.2f}s")
|
|
|
|
except optuna.TrialPruned:
|
|
# Propagate prune signal, objective will be set to worst later by Optuna
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Trial {trial.number}: Failed during classic run phase: {e}", exc_info=True)
|
|
validation_metric_value = worst_value # Assign worst value if classic run fails
|
|
finally:
|
|
# Clean up GPU memory after the run
|
|
del model_cl, trainer_cl, train_loader_cl, val_loader_cl, test_loader_cl
|
|
if torch.cuda.is_available(): torch.cuda.empty_cache()
|
|
|
|
|
|
# --- 4. Return Objective ---
|
|
logger.info(f"--- Trial {trial.number}: Finished ---")
|
|
logger.info(f" Objective (Validation {validation_metric_monitor}): {validation_metric_value:.5f}")
|
|
logger.info(f" Total time: {time.perf_counter() - trial_start_time:.2f}s")
|
|
|
|
# Return the single objective (validation metric)
|
|
return float(validation_metric_value)
|
|
|
|
|
|
# --- Main HPO Execution ---
|
|
def run_hpo():
|
|
"""Main execution function for HPO."""
|
|
args = parse_arguments()
|
|
config_path = Path(args.config)
|
|
try:
|
|
base_config = load_config(config_path) # Load base config once
|
|
logger.info(f"Successfully loaded configuration from {config_path}")
|
|
except Exception as e:
|
|
logger.critical(f"Failed to load configuration from {config_path}: {e}", exc_info=True)
|
|
sys.exit(1)
|
|
|
|
# Setup output dir...
|
|
if args.output_dir:
|
|
hpo_base_output_dir = Path(args.output_dir) # Use specific name for HPO dir
|
|
logger.info(f"Using HPO output directory from command line: {hpo_base_output_dir}")
|
|
elif base_config.optuna.storage and base_config.optuna.storage.startswith("sqlite:///"):
|
|
hpo_base_output_dir = Path(base_config.optuna.storage.replace("sqlite:///", "")).parent
|
|
logger.info(f"Using HPO output directory from Optuna storage path: {hpo_base_output_dir}")
|
|
else:
|
|
# Use output_dir from main config if available, otherwise default
|
|
main_output_dir = Path(getattr(base_config, 'output_dir', 'output'))
|
|
hpo_base_output_dir = main_output_dir / 'hpo_results'
|
|
logger.info(f"Using HPO output directory: {hpo_base_output_dir}")
|
|
hpo_base_output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Setup logging... (ensure file handler uses hpo_base_output_dir)
|
|
try:
|
|
level_name = base_config.log_level.upper()
|
|
effective_log_level = logging.getLevelName(level_name)
|
|
log_file = hpo_base_output_dir / f"{base_config.optuna.study_name}_hpo.log"
|
|
file_handler = logging.FileHandler(log_file, mode='a')
|
|
formatter = logging.Formatter('%(asctime)s - %(name)-25s - %(levelname)-7s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
|
|
file_handler.setFormatter(formatter)
|
|
# Add handler only if it's not already added (e.g., if run_hpo is called multiple times)
|
|
if not any(isinstance(h, logging.FileHandler) and h.baseFilename == str(log_file.resolve()) for h in logger.handlers):
|
|
logger.addHandler(file_handler)
|
|
logger.setLevel(effective_log_level)
|
|
logger.info(f"Set log level to {level_name}. Logging HPO run to console and {log_file}")
|
|
if effective_log_level <= logging.DEBUG: logger.debug("Debug logging enabled.")
|
|
except (AttributeError, ValueError) as e:
|
|
logger.warning(f"Could not set log level from config. Defaulting to INFO. Error: {e}")
|
|
logger.setLevel(logging.INFO)
|
|
|
|
|
|
# Setup seeding...
|
|
try:
|
|
seed = base_config.random_seed
|
|
set_seeds(seed)
|
|
except AttributeError:
|
|
logger.warning("Config missing 'random_seed'. Using default seed 42.")
|
|
set_seeds(42)
|
|
|
|
# Load data...
|
|
try:
|
|
logger.info("Loading base dataset for HPO...")
|
|
df = load_raw_data(base_config.data)
|
|
logger.info(f"Base dataset loaded. Shape: {df.shape}")
|
|
except Exception as e:
|
|
logger.critical(f"Failed to load raw data for HPO: {e}", exc_info=True)
|
|
sys.exit(1)
|
|
|
|
# --- Optuna Study Setup ---
|
|
try:
|
|
hpo_config = base_config.optuna
|
|
if not hpo_config.enabled:
|
|
logger.info("Optuna optimization is disabled in the configuration.")
|
|
sys.exit(0)
|
|
except AttributeError:
|
|
logger.critical("Optuna configuration section ('optuna') missing.")
|
|
sys.exit(1)
|
|
|
|
storage_path = hpo_config.storage
|
|
if storage_path and storage_path.startswith("sqlite:///"):
|
|
db_path = Path(storage_path.replace("sqlite:///", ""))
|
|
if not db_path.is_absolute():
|
|
db_path = hpo_base_output_dir / db_path # Relative to HPO output dir
|
|
db_path.parent.mkdir(parents=True, exist_ok=True)
|
|
storage_path = f"sqlite:///{db_path.resolve()}"
|
|
logger.info(f"Using Optuna storage: {storage_path}")
|
|
elif storage_path:
|
|
logger.info(f"Using Optuna storage: {storage_path} (non-SQLite)")
|
|
else:
|
|
logger.warning("No Optuna storage DB specified in config, using in-memory storage.")
|
|
|
|
try:
|
|
# Change to single objective 'minimize'
|
|
study = optuna.create_study(
|
|
study_name=hpo_config.study_name,
|
|
storage=storage_path,
|
|
direction="minimize", # Changed to single direction
|
|
load_if_exists=True,
|
|
pruner=optuna.pruners.MedianPruner() if hpo_config.pruning else optuna.pruners.NopPruner()
|
|
)
|
|
# Remove multi-objective check/attribute setting
|
|
# if not study._is_multi_objective:
|
|
# logger.warning(f"Study '{hpo_config.study_name}' exists but is not multi-objective.")
|
|
# elif 'objective_names' not in study.user_attrs:
|
|
# study.set_user_attr('objective_names', objective_names)
|
|
|
|
|
|
# --- Run Optimization ---
|
|
logger.info(f"Starting Optuna single-objective optimization: study='{hpo_config.study_name}', n_trials={hpo_config.n_trials}") # Updated log message
|
|
study.optimize(
|
|
lambda trial: objective(trial, base_config, df), # Pass base_config
|
|
n_trials=hpo_config.n_trials,
|
|
timeout=None,
|
|
gc_after_trial=True
|
|
)
|
|
|
|
# --- Report and Process Best Trial ---
|
|
logger.info("--- Optuna HPO Finished ---")
|
|
logger.info(f"Number of finished trials: {len(study.trials)}")
|
|
|
|
# Get the single best trial
|
|
best_trial = study.best_trial
|
|
if best_trial is None:
|
|
logger.warning("Optuna study finished, but no successful trial was completed.")
|
|
else:
|
|
logger.info(f"Best trial found (Trial {best_trial.number}):")
|
|
# Log details for the best trial
|
|
validation_metric_monitor = base_config.optuna.metric_to_optimize
|
|
logger.info(f" Objective ({validation_metric_monitor}): {best_trial.value:.5f}") # Use .value for single objective
|
|
logger.info(f" Hyperparameters:")
|
|
for key, value in best_trial.params.items():
|
|
logger.info(f" {key}: {value}")
|
|
|
|
# --- Re-run and Save Artifacts for the Best Trial ---
|
|
logger.info(f"-> Re-running Best Trial {best_trial.number} to save artifacts...")
|
|
trial_output_dir = hpo_base_output_dir / f"best_trial_num{best_trial.number}" # Simplified directory name
|
|
trial_output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
try:
|
|
# 1. Create config for this trial
|
|
best_config_dict = copy.deepcopy(base_config.model_dump(mode='python'))
|
|
# Update with trial's hyperparameters
|
|
for key, value in best_trial.params.items(): # Use best_trial.params
|
|
# Need to place the param in the correct nested config dict
|
|
if key in best_config_dict['training']: best_config_dict['training'][key] = value
|
|
elif key in best_config_dict['model']: best_config_dict['model'][key] = value
|
|
elif key in best_config_dict['features']: best_config_dict['features'][key] = value
|
|
# Add more sections if HPO tunes them
|
|
|
|
# Add non-tuned features forecast_horizon back
|
|
best_config_dict['features']['forecast_horizon'] = base_config.features.forecast_horizon
|
|
|
|
# Ensure evaluation plots and model saving are enabled for this final run
|
|
best_config_dict['evaluation']['save_plots'] = True
|
|
best_config_dict['training']['save_model'] = True # Assuming you want to save the best model
|
|
|
|
# Validate the final config for this trial
|
|
best_trial_config = MainConfig(**best_config_dict)
|
|
|
|
# Save the specific config used for this run
|
|
with open(trial_output_dir / "best_config.yaml", 'w') as f:
|
|
yaml.dump(best_config_dict, f, default_flow_style=False, sort_keys=False)
|
|
|
|
# 2. Run classic training (which saves model & plots)
|
|
logger.info(f"-> Running classic training for Best Trial {best_trial.number}...")
|
|
# Pass the specific config and output directory
|
|
run_classic_training(
|
|
config=best_trial_config,
|
|
full_df=df,
|
|
output_base_dir=trial_output_dir # outputs -> hpo_results/best_trial_num<n>/classic_run/
|
|
)
|
|
logger.info(f"-> Finished re-running and saving artifacts for Best Trial {best_trial.number} to {trial_output_dir}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"-> Failed to re-run or save artifacts for Best Trial {best_trial.number}: {e}", exc_info=True)
|
|
|
|
# --- Save Best Hyperparameters ---
|
|
best_params_file = hpo_base_output_dir / f"{hpo_config.study_name}_best_params.json" # Simplified filename
|
|
try:
|
|
with open(best_params_file, 'w') as f:
|
|
import json
|
|
json.dump(best_trial.params, f, indent=4) # Use best_trial.params
|
|
logger.info(f"Hyperparameters of the best trial saved to {best_params_file}")
|
|
except Exception as e:
|
|
logger.error(f"Failed to save parameters for best trial: {e}")
|
|
|
|
except Exception as e:
|
|
logger.critical(f"A critical error occurred during the Optuna study: {e}", exc_info=True)
|
|
sys.exit(1)
|
|
|
|
if __name__ == "__main__":
|
|
run_hpo() |