233 lines
8.8 KiB
Python
233 lines
8.8 KiB
Python
from __future__ import annotations
|
|
|
|
"""Optuna hyper-parameter search for EdgeNeXt on Naruto Sign (PROTOCOL §7-§9).
|
|
|
|
Single-objective (maximize val macro-F1) or multi-objective
|
|
(maximize macro-F1 + minimize trainable params) search with configurable
|
|
sampler and pruner. Define-by-run: the search space is described inside
|
|
``suggest_configs`` using ``trial.suggest_*`` calls, so conditional spaces
|
|
(e.g. focal-only ``focal_gamma``) are expressed naturally.
|
|
|
|
Usage::
|
|
|
|
python -m src.optuna_search --data-root /path/Naruto_Sign \
|
|
--n-trials 60 --sampler tpe --pruner median --study-name nss_v1
|
|
|
|
# Multi-objective (accuracy vs trainable params):
|
|
python -m src.optuna_search --data-root /path/Naruto_Sign \
|
|
--n-trials 80 --multi-objective --sampler nsga
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
|
|
import optuna
|
|
import torch
|
|
|
|
from .data import AugConfig, DataConfig
|
|
from .model import ModelConfig
|
|
from .train import TrainConfig, train_model
|
|
|
|
EDGENEXT_VARIANTS = [
|
|
"edgenext_xx_small.in1k",
|
|
"edgenext_x_small.in1k",
|
|
"edgenext_small.usi_in1k",
|
|
]
|
|
|
|
|
|
def suggest_configs(
|
|
trial: optuna.Trial, data_root: str, num_classes: int, epochs: int
|
|
) -> tuple[ModelConfig, DataConfig, TrainConfig]:
|
|
"""Sample one configuration from the search space (define-by-run).
|
|
|
|
Args:
|
|
trial: The Optuna trial.
|
|
data_root: Path to the ImageFolder dataset root.
|
|
num_classes: Number of classes.
|
|
epochs: Max epochs per trial (training budget).
|
|
|
|
Returns:
|
|
``(model_cfg, data_cfg, train_cfg)`` for :func:`train_model`.
|
|
"""
|
|
# --- backbone + transfer-learning regime (H2/H3) ----------------------
|
|
model_name = trial.suggest_categorical("model_name", EDGENEXT_VARIANTS)
|
|
regime = trial.suggest_categorical("regime", ["full", "partial", "mona"])
|
|
drop_path = trial.suggest_float("drop_path_rate", 0.0, 0.3)
|
|
drop_rate = trial.suggest_float("drop_rate", 0.0, 0.4)
|
|
|
|
n_unfrozen, mona_stages, mona_bottleneck, mona_kernels = 1, (2, 3), 64, (3, 5, 7)
|
|
if regime == "partial":
|
|
n_unfrozen = trial.suggest_int("n_unfrozen_stages", 0, 4)
|
|
elif regime == "mona":
|
|
mona_bottleneck = trial.suggest_categorical("mona_bottleneck", [16, 32, 64, 96])
|
|
kernel_set = trial.suggest_categorical(
|
|
"mona_kernels", ["3", "3-5", "3-5-7", "3-5-7-9"]
|
|
)
|
|
mona_kernels = tuple(int(k) for k in kernel_set.split("-"))
|
|
last_only = trial.suggest_categorical("mona_last_only", [True, False])
|
|
mona_stages = (3,) if last_only else (2, 3)
|
|
|
|
model_cfg = ModelConfig(
|
|
model_name=model_name,
|
|
num_classes=num_classes,
|
|
pretrained=True,
|
|
drop_rate=drop_rate,
|
|
drop_path_rate=drop_path,
|
|
regime=regime,
|
|
n_unfrozen_stages=n_unfrozen,
|
|
mona_stages=mona_stages,
|
|
mona_bottleneck=mona_bottleneck,
|
|
mona_kernels=mona_kernels,
|
|
)
|
|
|
|
# --- optimization (H1) -----------------------------------------------
|
|
lr = trial.suggest_float("lr", 1e-5, 5e-3, log=True)
|
|
weight_decay = trial.suggest_float("weight_decay", 1e-6, 1e-2, log=True)
|
|
optimizer = trial.suggest_categorical("optimizer", ["adamw", "sgd"])
|
|
backbone_lr_mult = trial.suggest_float("backbone_lr_mult", 0.02, 1.0, log=True)
|
|
|
|
# --- loss + imbalance handling (H4) ----------------------------------
|
|
loss_name = trial.suggest_categorical(
|
|
"loss_name", ["ce", "ce_ls", "weighted_ce", "focal"]
|
|
)
|
|
label_smoothing = trial.suggest_float("label_smoothing", 0.0, 0.2)
|
|
focal_gamma = trial.suggest_float("focal_gamma", 0.5, 4.0) if loss_name == "focal" else 2.0
|
|
use_weighted_sampler = trial.suggest_categorical("weighted_sampler", [True, False])
|
|
|
|
# --- augmentation (H6) -----------------------------------------------
|
|
img_size = trial.suggest_categorical("img_size", [224, 256])
|
|
mixup_alpha = trial.suggest_float("mixup_alpha", 0.0, 0.4)
|
|
rrc_scale_min = trial.suggest_float("rrc_scale_min", 0.5, 0.9)
|
|
randaug = trial.suggest_int("randaug_magnitude", 0, 9)
|
|
|
|
aug = AugConfig(
|
|
img_size=img_size,
|
|
use_hflip=False, # fixed OFF: chirality-sensitive seals (PROTOCOL §3)
|
|
rrc_scale_min=rrc_scale_min,
|
|
randaug_magnitude=randaug,
|
|
)
|
|
data_cfg = DataConfig(
|
|
data_root=data_root,
|
|
batch_size=trial.suggest_categorical("batch_size", [16, 32, 64]),
|
|
use_weighted_sampler=use_weighted_sampler,
|
|
aug=aug,
|
|
)
|
|
train_cfg = TrainConfig(
|
|
epochs=epochs,
|
|
lr=lr,
|
|
weight_decay=weight_decay,
|
|
optimizer=optimizer,
|
|
backbone_lr_mult=backbone_lr_mult,
|
|
loss_name=loss_name,
|
|
label_smoothing=label_smoothing,
|
|
focal_gamma=focal_gamma,
|
|
mixup_alpha=mixup_alpha,
|
|
)
|
|
return model_cfg, data_cfg, train_cfg
|
|
|
|
|
|
def make_study(args: argparse.Namespace) -> optuna.Study:
|
|
"""Construct a study with the requested sampler + pruner."""
|
|
samplers = {
|
|
"tpe": optuna.samplers.TPESampler(multivariate=True, seed=args.seed),
|
|
"random": optuna.samplers.RandomSampler(seed=args.seed),
|
|
"cmaes": optuna.samplers.CmaEsSampler(seed=args.seed),
|
|
"nsga": optuna.samplers.NSGAIISampler(seed=args.seed),
|
|
}
|
|
# AutoSampler lives in OptunaHub; load it when requested (PROTOCOL §2.5).
|
|
if args.sampler == "auto":
|
|
import optunahub
|
|
|
|
sampler = optunahub.load_module("samplers/auto_sampler").AutoSampler()
|
|
else:
|
|
sampler = samplers[args.sampler]
|
|
|
|
pruners = {
|
|
"median": optuna.pruners.MedianPruner(n_warmup_steps=5),
|
|
"asha": optuna.pruners.SuccessiveHalvingPruner(),
|
|
"hyperband": optuna.pruners.HyperbandPruner(),
|
|
"none": optuna.pruners.NopPruner(),
|
|
}
|
|
storage = f"sqlite:///{args.study_name}.db" if args.storage else None
|
|
|
|
if args.multi_objective:
|
|
return optuna.create_study(
|
|
study_name=args.study_name,
|
|
storage=storage,
|
|
load_if_exists=bool(storage),
|
|
sampler=sampler,
|
|
directions=["maximize", "minimize"], # macro-F1 up, params down
|
|
)
|
|
return optuna.create_study(
|
|
study_name=args.study_name,
|
|
storage=storage,
|
|
load_if_exists=bool(storage),
|
|
sampler=sampler,
|
|
pruner=pruners[args.pruner],
|
|
direction="maximize",
|
|
)
|
|
|
|
|
|
def main() -> None:
|
|
"""CLI entry point."""
|
|
p = argparse.ArgumentParser(description="Optuna HPO for EdgeNeXt / Naruto Sign")
|
|
p.add_argument("--data-root", required=True)
|
|
p.add_argument("--num-classes", type=int, default=13) # 12 seals + 'zero'
|
|
p.add_argument("--n-trials", type=int, default=60)
|
|
p.add_argument("--epochs", type=int, default=30)
|
|
p.add_argument("--sampler", choices=["tpe", "random", "cmaes", "nsga", "auto"], default="tpe")
|
|
p.add_argument("--pruner", choices=["median", "asha", "hyperband", "none"], default="median")
|
|
p.add_argument("--multi-objective", action="store_true")
|
|
p.add_argument("--study-name", default="nss_hpo")
|
|
p.add_argument("--storage", action="store_true", help="persist to sqlite db")
|
|
p.add_argument("--out", default="results")
|
|
p.add_argument("--seed", type=int, default=42)
|
|
args = p.parse_args()
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
os.makedirs(args.out, exist_ok=True)
|
|
|
|
def objective(trial: optuna.Trial):
|
|
model_cfg, data_cfg, train_cfg = suggest_configs(
|
|
trial, args.data_root, args.num_classes, args.epochs
|
|
)
|
|
# Multi-objective studies do not support pruning -> do not pass the trial
|
|
# (trial.report/should_prune raise on multi-objective trials).
|
|
prune_trial = None if args.multi_objective else trial
|
|
result = train_model(model_cfg, data_cfg, train_cfg, device, trial=prune_trial)
|
|
trial.set_user_attr("trainable_params", result["trainable_params"])
|
|
if args.multi_objective:
|
|
return result["best_val_macro_f1"], float(result["trainable_params"])
|
|
return result["best_val_macro_f1"]
|
|
|
|
study = make_study(args)
|
|
study.optimize(objective, n_trials=args.n_trials, gc_after_trial=True)
|
|
|
|
_persist(study, args)
|
|
|
|
|
|
def _persist(study: optuna.Study, args: argparse.Namespace) -> None:
|
|
"""Write best params / trials dataframe atomically."""
|
|
df = study.trials_dataframe()
|
|
df.to_csv(os.path.join(args.out, f"{args.study_name}_trials.csv"), index=False)
|
|
|
|
if args.multi_objective:
|
|
best = [
|
|
{"values": t.values, "params": t.params} for t in study.best_trials
|
|
]
|
|
payload = {"pareto_front": best}
|
|
else:
|
|
payload = {"best_value": study.best_value, "best_params": study.best_params}
|
|
|
|
tmp = os.path.join(args.out, f"{args.study_name}_best.json.tmp")
|
|
final = os.path.join(args.out, f"{args.study_name}_best.json")
|
|
with open(tmp, "w", encoding="utf-8") as f:
|
|
json.dump(payload, f, ensure_ascii=False, indent=2)
|
|
os.replace(tmp, final)
|
|
print(f"[optuna] wrote {final}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|