HPO-SH.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

from accelerate import notebook_launcher
from datetime import datetime
from torch.utils.data import DataLoader
from config.config import ImageConfig
from data.image.image import load_butterflies
from model.image.model import create_model
from model.image.train import train_loop, train_loop_SH
from components.components import create_lr_scheduler, create_noise_scheduler, create_optimizer
from evaluation.image.evaluate import evaluate

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--name", type=str, help="Name of study", required=True)
    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-3, log=True),
        "lr_scheduler": trial.suggest_categorical("lr_scheduler", ["constant", "constant_with_warmup", "linear", "cosine"]),
        "batch_size": trial.suggest_categorical("batch_size", [32, 64, 128]),
        "optimizer": trial.suggest_categorical("optimizer", ["Adam", "AdamW"]),
        "beta_scheduler": trial.suggest_categorical("beta_scheduler", ["linear", "scaled_linear", "squaredcos_cap_v2", "sigmoid"]),
    }
    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)

    hyperparameters["output_dir"] = output_dir
    config = ImageConfig(**hyperparameters)

    # = 2. Create a DataLoader =
    dataset = load_butterflies(config)
    dataloader = DataLoader(
        dataset,
        batch_size=config.batch_size,
        shuffle=True,
        num_workers=6,
        pin_memory=True,
        persistent_workers=False
    )

    # = 3. Calculate Warm-up Steps =
    config.lr_warmup_steps = config.lr_warmup_ratio * (config.nepochs * len(dataloader))

    # = 4. Prepare Model and Components =
    model = create_model(config)
    noise_scheduler = create_noise_scheduler(config)
    optimizer = create_optimizer(model, config)
    lr_scheduler = create_lr_scheduler(optimizer, config, len(dataloader))

    # = 5. Start Training =
    train_args = (config, model, noise_scheduler, optimizer, lr_scheduler, dataloader, trial)
    notebook_launcher(train_loop_SH, train_args, num_processes=1)

    # = 6. Evaluate Trial =
    metrics = evaluate(config, trial.number)

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

    return metrics["frechet_inception_distance"]


if __name__ == "__main__":

    # = 1. Retrieve Arguments =
    args = get_args()
    dataset_name = "butterflies64"
    config = ImageConfig()
    continue_study = args.resume
    global output_dir, generations_dir

    # = 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. 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}.")


    pruner = optuna.pruners.SuccessiveHalvingPruner(
        min_resource=6,                # min epochs for each trial
        reduction_factor=2              # keep best 1/2 fraction
    )

    study = optuna.create_study(
        study_name=study_name,
        direction="minimize",
        storage=storage_path,
        load_if_exists=True,
        sampler=optuna.samplers.RandomSampler(seed=34),
        pruner=pruner
    )

    # = 4. 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}).")

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