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()