HPO-2D.py

import sys
import os

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

import argparse
import optuna
import torch
import numpy as np
import time
import gc
import math

from accelerate import notebook_launcher
from datetime import datetime
from torch.utils.data import DataLoader
from config.config import TwoDConfig
from data.twodimensional.twodimensional import load_2d_data
from model.twodimensional.model import MLPModel
from model.twodimensional.train import train_loop
from components.components import create_lr_scheduler, create_optimizer
from model.twodimensional.sampling import set_diffusion_hyperparameters
from evaluation.twodimensional.evaluate import evaluate_model
from evaluation.twodimensional.visualize import plot_results


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", type=str, help="Type of dataset", choices=["swissroll", "moons", "spiral"], required=True)
    parser.add_argument("--name", type=str, help="Name of study", required=True)
    parser.add_argument("--sampler", type=str, help="Type of search strategy", required=True, choices=["RS", "TPE"])
    parser.add_argument("--resume", action="store_true", help="Resume a previous study")
    return parser.parse_args()

def objective(trial):

    # = 1. Create Search Space =
    hyperparameters = {
        "learning_rate": trial.suggest_float("learning_rate", 1e-5, 1e-2, log=True),
        "lr_scheduler": trial.suggest_categorical("lr_scheduler", ["constant", "constant_with_warmup", "linear", "cosine"]),
        "batch_size": trial.suggest_categorical("batch_size", [1024, 2048, 4096]),
        "optimizer": trial.suggest_categorical("optimizer", ["Adam", "AdamW"]),
        "diffusion_steps": trial.suggest_int("diffusion_steps", 20, 100, step=20),
    }
    if hyperparameters["optimizer"] == "AdamW":
        hyperparameters["weight_decay"] = trial.suggest_float("weight_decay", 1e-6, 1e-2, log=True)

    if hyperparameters["lr_scheduler"] != "constant":
        hyperparameters["lr_warmup_ratio"] = trial.suggest_float("lr_warmup_ratio", 0.0, 0.2, step=0.05)

    config = TwoDConfig(**hyperparameters)
    set_diffusion_hyperparameters(config)

    # = 2. Calculate Warm-up Steps =
    training_steps_epoch = (math.ceil(len(train_data) / config.batch_size))
    config.lr_warmup_steps = int(config.lr_warmup_ratio * (training_steps_epoch * config.nepochs))

    # = 3. Prepare Model and Components =
    model = MLPModel(nfeatures=config.nfeatures, nblocks=config.nblocks, nunits=config.nunits).to("cuda")
    optimizer = create_optimizer(model, config)
    lr_scheduler = create_lr_scheduler(optimizer, config, training_steps_epoch)

    # = 4. Start Training =
    train_loop(config, model, optimizer, lr_scheduler, train_data, eval_data)

    # = 5. Evaluate Trial =
    metric, generated_data = evaluate_model(model, config, eval_data)
    plot_results(train_data.cpu().numpy(), generated_data.numpy(), trial.number, metric, generations_dir)

    # = 6. Clean Components of Current Trial
    del model, optimizer, lr_scheduler
    torch.cuda.empty_cache()
    gc.collect()

    return metric


if __name__ == "__main__":

    # = 1. Retrieve Arguments =
    args = get_args()
    dataset_name = args.dataset
    config = TwoDConfig()
    continue_study = args.resume
    global output_dir, generations_dir, data_type
    data_type = dataset_name

    # = 2. Create Output Directory for Study =
    output_dir = os.path.join(f"studies/{dataset_name}", f"{args.name}")
    generations_dir = os.path.join(output_dir, "generations")

    # Check if directory already exists and we don't resume an existing study
    if os.path.exists(output_dir):
        if not continue_study:
            print(f"Output directory already exists ({output_dir}), use --resume to continue training.")
            sys.exit(1)
    
    else:
        # If continue, but no directory exists: exit
        if continue_study:
            print("Nothing to resume.")
            sys.exit(1)

        # Create directory
        os.makedirs(generations_dir, exist_ok=True)
        #os.makedirs(os.path.join(output_dir, "ddpm"))

    # = 3. Load Training and Evaluation Data =
    train_data, eval_data = load_2d_data(seed=34, data_type=data_type)

    # = 4. Create Study =
    study_name = f"{args.name}_study"
    storage_path = f"sqlite:///{os.path.join(output_dir, 'storage')}"
    print(f"Study name: {study_name}.")
    print(f"Storage path: {storage_path}.")

    if args.sampler == "TPE":
        sampler = optuna.samplers.TPESampler(seed=34)
    elif args.sampler == "RS":
        sampler = optuna.samplers.RandomSampler(seed=34)

    study = optuna.create_study(
        study_name=study_name,
        direction="minimize",
        storage=storage_path,
        load_if_exists=True,
        sampler=sampler
    )

    # = 5. Start the Study =
    start_time = time.time()
    try:
        study.optimize(objective, n_trials=50)
    except KeyboardInterrupt:
        print("Optimization cancelled.")
    finally:
        torch.cuda.empty_cache()
        gc.collect()
    total_time = time.time() - start_time

    print(f"\nBest parameters found: {study.best_params}.")
    print(f"Best score achieved: {study.best_value}.")
    print(f"Total study time: {total_time} seconds ({total_time/60:.2f}).")

    # = 6. Save Study to CSV File = 
    df = study.trials_dataframe()
    df.to_csv(os.path.join(output_dir, "results.csv"), index=False)