jung-geun/PSO
1
0
Fork 0
mirror of https://github.com/jung-geun/PSO.git synced 2026-09-20 14:11:48 +09:00
Code Issues Packages Projects Releases Wiki Activity
Files
master
PSO/test/tuning_suite.py
T
jung-geun 813433000a feat: modernize PSO and add convergence research 
Migrate the package and examples to the tensor-native PyTorch implementation, add benchmark evidence, and add the guarded post-training convergence protocol with TensorBoard progress monitoring and hash-verified recovery.

Constraint: Preserve one-shot official-test sealing and auditable research artifacts

Rejected: Commit local .omc runs and downloaded datasets | multi-gigabyte runtime state is machine-local

Confidence: high

Scope-risk: broad

Not-tested: Production CUDA run on pieroot-server
2026-09-07 22:03:25 +09:00

1433 lines
52 KiB
Python
Raw Permalink Blame History

"""
MNIST Tuning & Particle Scaling Study Suite (v4.0 Protocol 1.0.0)
Implements a reproducible three-phase study:
1. Search Phase: Inner validation split (2400 train / 600 val, stratified) from first 3000 MNIST training examples.
Fit PCA32 whitening on inner train only; transform inner val.
Evaluates candidate space across 5 movement methods (adaptive_moment, inertia, constriction, local_best, quantum) over 3 seeds (51-53).
Ranks validation accuracy descending, then validation loss ascending.
2. Confirmation Phase: Full 3000 training examples and 1000 untouched test examples.
Fit PCA32 whitening on full 3000 train only; transform 1000 test.
Evaluates top candidate per method over 5 seeds (61-65).
3. Scaling Phase: Full 3000 training / 1000 test set.
Evaluates selected adaptive_moment winner across particle counts 30, 60, 90, 120
under fixed-epoch (80) and fixed-budget (~2400 particle-epochs) regimens over 5 seeds (71-75).
"""
import argparse
import dataclasses
import hashlib
import json
import csv
import math
import os
import sys
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split
sys.path.insert(0, str(Path(__file__).resolve().parent))
from benchmark_suite import (
METHOD_STYLE,
calc_stats,
compute_data_fingerprint,
compute_model_fingerprint,
extract_plugin_metadata,
get_hardware_provenance,
get_method_style,
get_t_crit,
make_mnist_model,
resolve_execution_device,
save_json_atomic,
sync_device,
)
from pso import Optimizer, __version__ as pso_version
TUNING_PROTOCOL_VERSION = "1.0.0"
@dataclasses.dataclass
class CandidateConfig:
method: str
candidate_label: str
description: str
c0: float = 1.49618
c1: float = 1.49618
w_min: float = 0.7298
w_max: float = 0.7298
velocity_limit_ratio: Optional[float] = 0.025
mutation_swarm: float = 0.02
negative_swarm: float = 0.0
method_options: Dict[str, Any] = dataclasses.field(default_factory=dict)
def to_optimizer_kwargs(self, quick: bool = False) -> Dict[str, Any]:
fitness_sz = 50 if quick else 2000
kwargs: Dict[str, Any] = {
"method": self.method,
"velocity_limit_ratio": self.velocity_limit_ratio,
"mutation_swarm": self.mutation_swarm,
"negative_swarm": self.negative_swarm,
"particle_min": -3.0,
"particle_max": 3.0,
"boundary_strategy": "reflect",
"initialization": "model_noise",
"initial_position_noise": 0.05,
"evaluation": "fixed_subset",
"fitness_size": fitness_sz,
"convergence": "none",
"refinement": "none",
}
if self.method in ("adaptive_moment", "inertia", "local_best"):
kwargs["c0"] = self.c0
kwargs["c1"] = self.c1
kwargs["w_min"] = self.w_min
kwargs["w_max"] = self.w_max
elif self.method == "constriction":
kwargs["c0"] = self.c0
kwargs["c1"] = self.c1
if self.method_options:
kwargs["method_options"] = dict(self.method_options)
return kwargs
def get_mnist_raw_data() -> Tuple[np.ndarray, np.ndarray, torch.Tensor, torch.Tensor]:
from torchvision.datasets import MNIST
cache_dir = Path("result/cache")
cache_dir.mkdir(parents=True, exist_ok=True)
train_dataset = MNIST(root=str(cache_dir), train=True, download=True)
test_dataset = MNIST(root=str(cache_dir), train=False, download=True)
x_train_raw = (train_dataset.data[:3000].float() / 255.0).reshape(3000, -1).numpy()
y_train = train_dataset.targets[:3000].long()
x_test_raw = (test_dataset.data[:1000].float() / 255.0).reshape(1000, -1).numpy()
y_test = test_dataset.targets[:1000].long()
return x_train_raw, x_test_raw, y_train, y_test
def get_search_candidates() -> Dict[str, List[CandidateConfig]]:
candidates: Dict[str, List[CandidateConfig]] = {
"adaptive_moment": [],
"inertia": [],
"constriction": [],
"local_best": [],
"quantum": [],
}
# --- 1. Adaptive Moment Candidates ---
blends = [0.03, 0.06, 0.10, 0.15]
steps = [0.5, 1.0, 1.5]
for b in blends:
for s in steps:
label = f"am_b{b}_s{s}"
candidates["adaptive_moment"].append(
CandidateConfig(
method="adaptive_moment",
candidate_label=label,
description=f"Adaptive Moment blend={b}, step={s}, beta1=0.9",
c0=1.49618,
c1=1.49618,
w_min=0.7298,
w_max=0.7298,
velocity_limit_ratio=0.025,
mutation_swarm=0.02,
method_options={"moment_blend": b, "moment_step_size": s, "moment_beta1": 0.9},
)
)
for beta1 in [0.8, 0.95]:
label = f"am_b0.06_s1.0_beta{beta1}"
candidates["adaptive_moment"].append(
CandidateConfig(
method="adaptive_moment",
candidate_label=label,
description=f"Adaptive Moment blend=0.06, step=1.0, beta1={beta1}",
c0=1.49618,
c1=1.49618,
w_min=0.7298,
w_max=0.7298,
velocity_limit_ratio=0.025,
mutation_swarm=0.02,
method_options={"moment_blend": 0.06, "moment_step_size": 1.0, "moment_beta1": beta1},
)
)
# --- 2. Inertia Candidates ---
candidates["inertia"] = [
CandidateConfig(
method="inertia",
candidate_label="inertia_canonical",
description="Canonical Inertia (c0=c1=2.0, w=0.9->0.4, vel=0.1, mut=0)",
c0=2.0,
c1=2.0,
w_min=0.4,
w_max=0.9,
velocity_limit_ratio=0.1,
mutation_swarm=0.0,
),
CandidateConfig(
method="inertia",
candidate_label="inertia_tuned",
description="Tuned Inertia (c0=c1=1.49618, w=0.7298, vel=0.025, mut=0.02)",
c0=1.49618,
c1=1.49618,
w_min=0.7298,
w_max=0.7298,
velocity_limit_ratio=0.025,
mutation_swarm=0.02,
),
CandidateConfig(
method="inertia",
candidate_label="inertia_low_w",
description="Low Inertia w (c0=c1=1.49618, w=0.55, vel=0.025, mut=0.02)",
c0=1.49618,
c1=1.49618,
w_min=0.55,
w_max=0.55,
velocity_limit_ratio=0.025,
mutation_swarm=0.02,
),
CandidateConfig(
method="inertia",
candidate_label="inertia_w_decay",
description="Decaying Inertia (c0=c1=1.49618, w=0.9->0.4, vel=0.025, mut=0.02)",
c0=1.49618,
c1=1.49618,
w_min=0.4,
w_max=0.9,
velocity_limit_ratio=0.025,
mutation_swarm=0.02,
),
CandidateConfig(
method="inertia",
candidate_label="inertia_asymmetric",
description="Asymmetric Inertia (c0=1.8, c1=1.2, w=0.7298, vel=0.025, mut=0.02)",
c0=1.8,
c1=1.2,
w_min=0.7298,
w_max=0.7298,
velocity_limit_ratio=0.025,
mutation_swarm=0.02,
),
]
# --- 3. Constriction Candidates ---
candidates["constriction"] = [
CandidateConfig(
method="constriction",
candidate_label="constriction_c201",
description="Constriction c0=c1=2.01 (vel=0.025, mut=0)",
c0=2.01,
c1=2.01,
w_min=0.7298,
w_max=0.7298,
velocity_limit_ratio=0.025,
mutation_swarm=0.0,
),
CandidateConfig(
method="constriction",
candidate_label="constriction_c205_canonical",
description="Constriction c0=c1=2.05 canonical-ish (vel=0.05, mut=0.02)",
c0=2.05,
c1=2.05,
w_min=0.7298,
w_max=0.7298,
velocity_limit_ratio=0.05,
mutation_swarm=0.02,
),
CandidateConfig(
method="constriction",
candidate_label="constriction_c205_tuned",
description="Constriction c0=c1=2.05 tuned (vel=0.025, mut=0)",
c0=2.05,
c1=2.05,
w_min=0.7298,
w_max=0.7298,
velocity_limit_ratio=0.025,
mutation_swarm=0.0,
),
CandidateConfig(
method="constriction",
candidate_label="constriction_c250",
description="Constriction c0=c1=2.50 (vel=0.025, mut=0)",
c0=2.50,
c1=2.50,
w_min=0.7298,
w_max=0.7298,
velocity_limit_ratio=0.025,
mutation_swarm=0.0,
),
CandidateConfig(
method="constriction",
candidate_label="constriction_asymmetric",
description="Asymmetric Constriction c0=2.8, c1=1.3 (vel=0.025, mut=0)",
c0=2.8,
c1=1.3,
w_min=0.7298,
w_max=0.7298,
velocity_limit_ratio=0.025,
mutation_swarm=0.0,
),
]
# --- 4. Local-Best Candidates ---
candidates["local_best"] = [
CandidateConfig(
method="local_best",
candidate_label="local_best_r1_constant",
description="Local Best Ring Radius 1 (constant w=0.7298)",
c0=1.49618,
c1=1.49618,
w_min=0.7298,
w_max=0.7298,
velocity_limit_ratio=0.025,
mutation_swarm=0.02,
method_options={"neighborhood_radius": 1, "c0": 1.49618, "c1": 1.49618, "w_min": 0.7298, "w_max": 0.7298},
),
CandidateConfig(
method="local_best",
candidate_label="local_best_r2_constant",
description="Local Best Ring Radius 2 (constant w=0.7298)",
c0=1.49618,
c1=1.49618,
w_min=0.7298,
w_max=0.7298,
velocity_limit_ratio=0.025,
mutation_swarm=0.02,
method_options={"neighborhood_radius": 2, "c0": 1.49618, "c1": 1.49618, "w_min": 0.7298, "w_max": 0.7298},
),
CandidateConfig(
method="local_best",
candidate_label="local_best_r4_constant",
description="Local Best Ring Radius 4 (constant w=0.7298)",
c0=1.49618,
c1=1.49618,
w_min=0.7298,
w_max=0.7298,
velocity_limit_ratio=0.025,
mutation_swarm=0.02,
method_options={"neighborhood_radius": 4, "c0": 1.49618, "c1": 1.49618, "w_min": 0.7298, "w_max": 0.7298},
),
CandidateConfig(
method="local_best",
candidate_label="local_best_r1_decay",
description="Local Best Ring Radius 1 (decaying w=0.9->0.4)",
c0=1.49618,
c1=1.49618,
w_min=0.4,
w_max=0.9,
velocity_limit_ratio=0.025,
mutation_swarm=0.02,
method_options={"neighborhood_radius": 1, "c0": 1.49618, "c1": 1.49618, "w_min": 0.4, "w_max": 0.9},
),
]
# --- 5. Quantum Candidates ---
candidates["quantum"] = [
CandidateConfig(
method="quantum",
candidate_label="quantum_beta_0.5_1.0",
description="Quantum PSO (beta=1.0->0.5)",
c0=1.49618,
c1=1.49618,
w_min=0.7298,
w_max=0.7298,
velocity_limit_ratio=None,
mutation_swarm=0.0,
negative_swarm=0.0,
method_options={"beta_min": 0.5, "beta_max": 1.0},
),
CandidateConfig(
method="quantum",
candidate_label="quantum_beta_0.6_1.0",
description="Quantum PSO (beta=1.0->0.6)",
c0=1.49618,
c1=1.49618,
w_min=0.7298,
w_max=0.7298,
velocity_limit_ratio=None,
mutation_swarm=0.0,
negative_swarm=0.0,
method_options={"beta_min": 0.6, "beta_max": 1.0},
),
CandidateConfig(
method="quantum",
candidate_label="quantum_beta_0.5_1.2",
description="Quantum PSO (beta=1.2->0.5)",
c0=1.49618,
c1=1.49618,
w_min=0.7298,
w_max=0.7298,
velocity_limit_ratio=None,
mutation_swarm=0.0,
negative_swarm=0.0,
method_options={"beta_min": 0.5, "beta_max": 1.2},
),
CandidateConfig(
method="quantum",
candidate_label="quantum_beta_0.4_0.9",
description="Quantum PSO (beta=0.9->0.4)",
c0=1.49618,
c1=1.49618,
w_min=0.7298,
w_max=0.7298,
velocity_limit_ratio=None,
mutation_swarm=0.0,
negative_swarm=0.0,
method_options={"beta_min": 0.4, "beta_max": 0.9},
),
]
return candidates
def run_single_experiment(
cfg: CandidateConfig,
seed: int,
x_train: torch.Tensor,
y_train: torch.Tensor,
x_eval: torch.Tensor,
y_eval: torch.Tensor,
n_particles: int,
epochs: int,
batch_size: int,
device: torch.device,
quick: bool = False,
eval_metric_name: str = "val",
data_fp: str = "",
run_type: str = "search",
extra_meta: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
opt_kwargs = cfg.to_optimizer_kwargs(quick=quick)
hw_provenance = get_hardware_provenance(device)
warmup_ep = min(2, epochs)
# --- Untimed Warmup Phase ---
warmup_model = make_mnist_model(seed=seed)
warmup_loss = nn.CrossEntropyLoss()
warmup_opt = Optimizer(
model=warmup_model,
loss=warmup_loss,
task="multiclass",
n_particles=n_particles,
seed=seed,
device=device,
**opt_kwargs,
)
warmup_opt.fit(
x_train,
y_train,
epochs=warmup_ep,
batch_size=batch_size,
renewal="loss",
)
sync_device(device)
del warmup_opt, warmup_model, warmup_loss
# --- Timed Fit Phase ---
model = make_mnist_model(seed=seed)
model_fp = compute_model_fingerprint(model)
loss_inst = nn.CrossEntropyLoss()
opt = Optimizer(
model=model,
loss=loss_inst,
task="multiclass",
n_particles=n_particles,
seed=seed,
device=device,
**opt_kwargs,
)
plugin_meta = extract_plugin_metadata(opt)
sync_device(device)
t0 = time.perf_counter()
train_loss, train_acc, train_mse = opt.fit(
x_train,
y_train,
epochs=epochs,
batch_size=batch_size,
renewal="loss",
)
sync_device(device)
t1 = time.perf_counter()
fit_time = t1 - t0
# Separate Evaluation on Validation or Test set
eval_loss, eval_acc, eval_mse = opt.evaluate(x_eval, y_eval)
full_resolved_config = dict(opt_kwargs)
full_resolved_config.update({
"n_particles": n_particles,
"epochs": epochs,
"batch_size": batch_size,
"renewal": "loss",
})
res = {
"protocol_version": TUNING_PROTOCOL_VERSION,
"pso_version": pso_version,
"torch_version": torch.__version__,
"hardware": hw_provenance,
"timing_scope": "fit_only_after_method_specific_warmup",
"warmup_epochs": warmup_ep,
"error": None,
"phase": run_type,
"method": cfg.method,
"candidate_label": cfg.candidate_label,
"seed": seed,
"n_particles": n_particles,
"epochs": epochs,
"particle_epochs": n_particles * epochs,
"train_loss": float(train_loss),
"train_acc": float(train_acc),
"train_mse": float(train_mse),
f"{eval_metric_name}_loss": float(eval_loss),
f"{eval_metric_name}_acc": float(eval_acc),
f"{eval_metric_name}_mse": float(eval_mse),
"fit_time_sec": float(fit_time),
"data_fingerprint": data_fp,
"model_fingerprint": model_fp,
"device": str(device),
"completed": True,
"plugins": plugin_meta,
"config": full_resolved_config,
}
if extra_meta:
res.update(extra_meta)
return res
def write_tuning_csvs(
search_runs: List[Dict[str, Any]],
confirmation_runs: List[Dict[str, Any]],
scaling_runs: List[Dict[str, Any]],
result_dir: Path,
):
result_dir.mkdir(parents=True, exist_ok=True)
# 1. Search CSV
search_csv = result_dir / "pso_v4_tuning_search.csv"
search_fields = [
"method", "candidate_label", "seed", "n_particles", "epochs", "particle_epochs",
"train_loss", "train_acc", "val_loss", "val_acc", "val_mse",
"fit_time_sec", "data_fingerprint", "model_fingerprint", "device"
]
with open(search_csv, "w", newline="", encoding="utf-8") as f:
writer = csv.DictWriter(f, fieldnames=search_fields, extrasaction="ignore")
writer.writeheader()
for r in search_runs:
if r.get("completed"):
writer.writerow(r)
# 2. Confirmation CSV
confirm_csv = result_dir / "pso_v4_tuning_confirmation.csv"
confirm_fields = [
"method", "candidate_label", "seed", "n_particles", "epochs", "particle_epochs",
"train_loss", "train_acc", "test_loss", "test_acc", "test_mse",
"fit_time_sec", "data_fingerprint", "model_fingerprint", "device"
]
with open(confirm_csv, "w", newline="", encoding="utf-8") as f:
writer = csv.DictWriter(f, fieldnames=confirm_fields, extrasaction="ignore")
writer.writeheader()
for r in confirmation_runs:
if r.get("completed"):
writer.writerow(r)
# 3. Particle Scaling CSV
scaling_csv = result_dir / "pso_v4_particle_scaling.csv"
scaling_fields = [
"method", "candidate_label", "regimen", "n_particles", "epochs", "particle_epochs", "seed",
"train_loss", "train_acc", "test_loss", "test_acc", "test_mse",
"fit_time_sec", "data_fingerprint", "model_fingerprint", "device"
]
with open(scaling_csv, "w", newline="", encoding="utf-8") as f:
writer = csv.DictWriter(f, fieldnames=scaling_fields, extrasaction="ignore")
writer.writeheader()
for r in scaling_runs:
if r.get("completed"):
writer.writerow(r)
def _persist_json_state(
output_json: Path,
device: torch.device,
quick: bool,
hw_provenance: Dict[str, Any],
split_fingerprints: Dict[str, str],
pca_provenance: Dict[str, Any],
all_search_candidates: Dict[str, List[CandidateConfig]],
search_runs: List[Dict[str, Any]],
confirmation_runs: List[Dict[str, Any]],
scaling_runs: List[Dict[str, Any]],
):
search_summaries: Dict[str, Any] = {}
for r in search_runs:
if not r.get("completed"):
continue
lbl = r["candidate_label"]
if lbl not in search_summaries:
search_summaries[lbl] = {
"method": r["method"],
"candidate_label": lbl,
"val_accs": [],
"val_losses": [],
"fit_times": [],
}
search_summaries[lbl]["val_accs"].append(r["val_acc"])
search_summaries[lbl]["val_losses"].append(r["val_loss"])
search_summaries[lbl]["fit_times"].append(r["fit_time_sec"])
for lbl, s in search_summaries.items():
s["val_acc_stats"] = calc_stats(s["val_accs"])
s["val_loss_stats"] = calc_stats(s["val_losses"])
s["fit_time_stats"] = calc_stats(s["fit_times"])
confirm_summaries: Dict[str, Any] = {}
for r in confirmation_runs:
if not r.get("completed"):
continue
m = r["method"]
if m not in confirm_summaries:
confirm_summaries[m] = {
"method": m,
"candidate_label": r["candidate_label"],
"test_accs": [],
"test_losses": [],
"fit_times": [],
}
confirm_summaries[m]["test_accs"].append(r["test_acc"])
confirm_summaries[m]["test_losses"].append(r["test_loss"])
confirm_summaries[m]["fit_times"].append(r["fit_time_sec"])
for m, s in confirm_summaries.items():
s["test_acc_stats"] = calc_stats(s["test_accs"])
s["test_loss_stats"] = calc_stats(s["test_losses"])
s["fit_time_stats"] = calc_stats(s["fit_times"])
scaling_summaries: Dict[str, Any] = {}
for r in scaling_runs:
if not r.get("completed"):
continue
key = f"{r['n_particles']}p_{r['epochs']}e_{r.get('regimen', 'unknown')}"
if key not in scaling_summaries:
scaling_summaries[key] = {
"n_particles": r["n_particles"],
"epochs": r["epochs"],
"particle_epochs": r["particle_epochs"],
"regimen": r.get("regimen", "unknown"),
"test_accs": [],
"test_losses": [],
"fit_times": [],
}
scaling_summaries[key]["test_accs"].append(r["test_acc"])
scaling_summaries[key]["test_losses"].append(r["test_loss"])
scaling_summaries[key]["fit_times"].append(r["fit_time_sec"])
for key, s in scaling_summaries.items():
s["test_acc_stats"] = calc_stats(s["test_accs"])
s["test_loss_stats"] = calc_stats(s["test_losses"])
s["fit_time_stats"] = calc_stats(s["fit_times"])
required_seeds = 1 if quick else 3
winners = {}
for method, cand_list in all_search_candidates.items():
cand_scores = []
for cfg in cand_list:
lbl = cfg.candidate_label
matching_runs = [r for r in search_runs if r.get("candidate_label") == lbl and r.get("completed")]
if len(matching_runs) < required_seeds:
continue
val_accs = [r["val_acc"] for r in matching_runs]
val_losses = [r["val_loss"] for r in matching_runs]
cand_scores.append({
"candidate_label": lbl,
"mean_val_loss": float(np.mean(val_losses)),
"mean_val_acc": float(np.mean(val_accs)),
"config": cfg.to_optimizer_kwargs(quick=quick),
})
if cand_scores:
# Rank validation accuracy descending, then validation loss ascending
cand_scores.sort(key=lambda c: (-c["mean_val_acc"], c["mean_val_loss"]))
winners[method] = cand_scores[0]
payload = {
"tuning_protocol_version": TUNING_PROTOCOL_VERSION,
"pso_version": pso_version,
"torch_version": torch.__version__,
"quick": quick,
"device": str(device),
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
"hardware": hw_provenance,
"split_fingerprints": split_fingerprints,
"pca_provenance": pca_provenance,
"selection_criteria": "Validation accuracy descending, then validation loss ascending across required search seeds",
"winners": winners,
"summaries": {
"search": search_summaries,
"confirmation": confirm_summaries,
"scaling": scaling_summaries,
},
"search_runs": search_runs,
"confirmation_runs": confirmation_runs,
"scaling_runs": scaling_runs,
}
save_json_atomic(payload, output_json)
def render_tuning_plots(
search_runs: List[Dict[str, Any]],
confirmation_runs: List[Dict[str, Any]],
scaling_runs: List[Dict[str, Any]],
winners: Dict[str, Dict[str, Any]],
figure_dir: Path,
):
figure_dir.mkdir(parents=True, exist_ok=True)
# -------------------------------------------------------------
# Figure 1: pso_v4_extended_tuning.png
# -------------------------------------------------------------
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5.5))
methods = ["adaptive_moment", "inertia", "constriction", "local_best", "quantum"]
method_labels = {
"adaptive_moment": "Adaptive Moment",
"inertia": "Inertia Weight",
"constriction": "Constriction",
"local_best": "Local Best",
"quantum": "Quantum PSO",
}
positions = []
box_data = []
winner_x = []
winner_y = []
x_ticks = []
x_tick_labels = []
for idx, m in enumerate(methods):
m_runs = [r for r in search_runs if r.get("method") == m and r.get("completed")]
if not m_runs:
continue
cand_means = {}
for r in m_runs:
lbl = r["candidate_label"]
if lbl not in cand_means:
cand_means[lbl] = []
cand_means[lbl].append(r["val_loss"])
c_means = [float(np.mean(vals)) for vals in cand_means.values()]
box_data.append(c_means)
pos = idx + 1
positions.append(pos)
x_ticks.append(pos)
x_tick_labels.append(method_labels.get(m, m))
win_info = winners.get(m)
if win_info and win_info["candidate_label"] in cand_means:
win_loss = float(np.mean(cand_means[win_info["candidate_label"]]))
winner_x.append(pos)
winner_y.append(win_loss)
if box_data:
bp = ax1.boxplot(
box_data,
positions=positions,
widths=0.45,
patch_artist=True,
showmeans=False,
)
for box, m in zip(bp["boxes"], methods[:len(box_data)]):
c, _ = get_method_style(m)
box.set_facecolor(c)
box.set_alpha(0.6)
box.set_edgecolor("#333333")
if winner_x:
ax1.scatter(
winner_x,
winner_y,
color="#D55E00",
marker="*",
s=180,
zorder=5,
label="Selected Winner Candidate",
)
ax1.set_xticks(x_ticks)
ax1.set_xticklabels(x_tick_labels, rotation=15, ha="right", fontsize=10)
ax1.set_ylabel("Validation Cross-Entropy Loss", fontsize=11)
ax1.set_title("Phase 1: Inner Validation Search (Candidates per Method)", fontsize=12, fontweight="bold")
ax1.grid(True, linestyle="--", alpha=0.5)
if winner_x:
ax1.legend(loc="upper right")
conf_x = []
conf_y = []
conf_ci = []
conf_colors = []
for idx, m in enumerate(methods):
m_runs = [r for r in confirmation_runs if r.get("method") == m and r.get("completed")]
if not m_runs:
continue
accs = [r["test_acc"] * 100.0 for r in m_runs]
stats = calc_stats(accs)
conf_x.append(idx + 1)
conf_y.append(stats["mean"])
conf_ci.append(stats["ci95_t"])
c, _ = get_method_style(m)
conf_colors.append(c)
if conf_x:
bars = ax2.bar(
conf_x,
conf_y,
yerr=conf_ci,
capsize=5,
color=conf_colors,
edgecolor="#333333",
alpha=0.85,
width=0.5,
)
ax2.set_xticks(conf_x)
ax2.set_xticklabels([method_labels.get(m, m) for m in methods[:len(conf_x)]], rotation=15, ha="right", fontsize=10)
ax2.set_ylabel("Held-Out Test Accuracy (%)", fontsize=11)
ax2.set_title("Phase 2: Held-Out Test Confirmation (Method Winners)", fontsize=12, fontweight="bold")
ax2.grid(True, linestyle="--", alpha=0.5)
for bar, y_val, ci_val in zip(bars, conf_y, conf_ci):
ax2.text(
bar.get_x() + bar.get_width() / 2.0,
y_val + ci_val + 0.5,
f"{y_val:.1f}%",
ha="center",
va="bottom",
fontsize=9,
fontweight="bold",
)
plt.tight_layout()
fig_path1 = figure_dir / "pso_v4_extended_tuning.png"
plt.savefig(fig_path1, dpi=300, bbox_inches="tight")
plt.close(fig)
print(f"Rendered plot: {fig_path1}")
# -------------------------------------------------------------
# Figure 2: pso_v4_particle_scaling.png
# -------------------------------------------------------------
fig, (ax_acc, ax_loss, ax_time) = plt.subplots(1, 3, figsize=(16, 4.8))
regimens = ["fixed_epoch", "fixed_budget"]
regimen_names = {
"fixed_epoch": "Fixed epochs: 80",
"fixed_budget": "Fixed budget: ~2,400 particle-epochs",
}
regimen_colors = {
"fixed_epoch": "#0072B2",
"fixed_budget": "#D55E00",
}
regimen_markers = {
"fixed_epoch": "o",
"fixed_budget": "X",
}
for reg in regimens:
reg_runs = [r for r in scaling_runs if r.get("regimen") == reg and r.get("completed")]
if not reg_runs:
continue
by_p: Dict[int, List[Dict[str, Any]]] = {}
for r in reg_runs:
p = r["n_particles"]
if p not in by_p:
by_p[p] = []
by_p[p].append(r)
p_sorted = sorted(by_p.keys())
x_offset = -1.2 if reg == "fixed_epoch" else 1.2
plot_x = [p + x_offset for p in p_sorted]
acc_means, acc_cis = [], []
loss_means, loss_cis = [], []
time_means, time_cis = [], []
for p in p_sorted:
p_runs = by_p[p]
acc_st = calc_stats([r["test_acc"] * 100.0 for r in p_runs])
loss_st = calc_stats([r["test_loss"] for r in p_runs])
time_st = calc_stats([r["fit_time_sec"] for r in p_runs])
acc_means.append(acc_st["mean"])
acc_cis.append(acc_st["ci95_t"])
loss_means.append(loss_st["mean"])
loss_cis.append(loss_st["ci95_t"])
time_means.append(time_st["mean"])
time_cis.append(time_st["ci95_t"])
color = regimen_colors[reg]
marker = regimen_markers[reg]
label = regimen_names[reg]
ax_acc.errorbar(
plot_x,
acc_means,
yerr=acc_cis,
fmt=f"-{marker}",
color=color,
linewidth=2,
markersize=6,
capsize=4,
label=label,
)
ax_loss.errorbar(
plot_x,
loss_means,
yerr=loss_cis,
fmt=f"-{marker}",
color=color,
linewidth=2,
markersize=6,
capsize=4,
label=label,
)
ax_time.errorbar(
plot_x,
time_means,
yerr=time_cis,
fmt=f"-{marker}",
color=color,
linewidth=2,
markersize=6,
capsize=4,
label=label,
)
ax_acc.set_title("Test Accuracy vs Particle Count", fontsize=11, fontweight="bold")
ax_acc.set_xlabel("Particle Count", fontsize=10)
ax_acc.set_ylabel("Test Accuracy (%)", fontsize=10)
ax_acc.set_xticks([30, 60, 90, 120])
ax_acc.grid(True, linestyle="--", alpha=0.5)
ax_loss.set_title("Test Cross-Entropy Loss vs Particle Count", fontsize=11, fontweight="bold")
ax_loss.set_xlabel("Particle Count", fontsize=10)
ax_loss.set_ylabel("Test Cross-Entropy Loss", fontsize=10)
ax_loss.set_xticks([30, 60, 90, 120])
ax_loss.grid(True, linestyle="--", alpha=0.5)
ax_time.set_title("Fit Runtime vs Particle Count", fontsize=11, fontweight="bold")
ax_time.set_xlabel("Particle Count", fontsize=10)
ax_time.set_ylabel("Fit Runtime (seconds)", fontsize=10)
ax_time.set_xticks([30, 60, 90, 120])
ax_time.grid(True, linestyle="--", alpha=0.5)
handles, labels = ax_acc.get_legend_handles_labels()
if handles:
fig.legend(
handles,
labels,
loc="upper center",
bbox_to_anchor=(0.5, 1.02),
ncol=2,
frameon=True,
)
plt.tight_layout(rect=(0.0, 0.0, 1.0, 0.91))
fig_path2 = figure_dir / "pso_v4_particle_scaling.png"
plt.savefig(fig_path2, dpi=300, bbox_inches="tight")
plt.close(fig)
print(f"Rendered plot: {fig_path2}")
def run_tuning_study(
phase: str = "all",
device_name: Optional[str] = None,
quick: bool = False,
method_filter: Optional[List[str]] = None,
overwrite: bool = False,
output_json: Path = Path("benchmark_results/pso_v4_tuning.json"),
result_dir: Path = Path("benchmark_results"),
figure_dir: Path = Path("history_plt"),
):
device = resolve_execution_device(device_name)
print(f"Executing MNIST Tuning Study on device: {device}")
hw_provenance = get_hardware_provenance(device)
existing_data: Dict[str, Any] = {}
completed_search_runs: Dict[str, Dict[str, Any]] = {}
completed_confirm_runs: Dict[str, Dict[str, Any]] = {}
completed_scaling_runs: Dict[str, Dict[str, Any]] = {}
if output_json.exists() and not overwrite:
try:
with open(output_json, "r", encoding="utf-8") as f:
existing_data = json.load(f)
if (
existing_data.get("tuning_protocol_version") == TUNING_PROTOCOL_VERSION
and existing_data.get("quick") == quick
and existing_data.get("device") == str(device)
):
for r in existing_data.get("search_runs", []):
if r.get("completed") and "run_id" in r:
completed_search_runs[r["run_id"]] = r
for r in existing_data.get("confirmation_runs", []):
if r.get("completed") and "run_id" in r:
completed_confirm_runs[r["run_id"]] = r
for r in existing_data.get("scaling_runs", []):
if r.get("completed") and "run_id" in r:
completed_scaling_runs[r["run_id"]] = r
print(
f"Loaded existing runs from {output_json}: "
f"{len(completed_search_runs)} search, "
f"{len(completed_confirm_runs)} confirmation, "
f"{len(completed_scaling_runs)} scaling."
)
else:
print("Existing JSON protocol version, quick mode, or device differs. Starting fresh.")
except Exception as e:
print(f"Warning: Failed to load existing JSON ({e}). Starting fresh.")
# Data loading and PCA preprocessing
x_train_raw, x_test_raw, y_train_3000, y_test_1000 = get_mnist_raw_data()
# Search inner split: Stratified 2400 train / 600 validation
x_inner_tr_raw, x_inner_val_raw, y_inner_tr_np, y_inner_val_np = train_test_split(
x_train_raw,
y_train_3000.numpy(),
train_size=2400,
test_size=600,
stratify=y_train_3000.numpy(),
random_state=42,
)
y_inner_tr = torch.tensor(y_inner_tr_np, dtype=torch.long)
y_inner_val = torch.tensor(y_inner_val_np, dtype=torch.long)
# Fit PCA32 whitening on inner train ONLY for Search phase
pca_search = PCA(n_components=32, whiten=True, random_state=42)
x_inner_tr = torch.tensor(pca_search.fit_transform(x_inner_tr_raw), dtype=torch.float32)
x_inner_val = torch.tensor(pca_search.transform(x_inner_val_raw), dtype=torch.float32)
# Fit PCA32 whitening on full 3000 train ONLY for Confirmation & Scaling phases
pca_full = PCA(n_components=32, whiten=True, random_state=42)
x_full_tr = torch.tensor(pca_full.fit_transform(x_train_raw), dtype=torch.float32)
x_full_test = torch.tensor(pca_full.transform(x_test_raw), dtype=torch.float32)
search_data_fp = compute_data_fingerprint(x_inner_tr, x_inner_val, y_inner_tr, y_inner_val)
full_data_fp = compute_data_fingerprint(x_full_tr, x_full_test, y_train_3000, y_test_1000)
split_fingerprints = {
"search_inner": search_data_fp,
"full": full_data_fp,
}
pca_provenance = {
"search": {
"n_samples_fit": 2400,
"n_samples_val": 600,
"n_components": 32,
"whiten": True,
"random_state": 42,
"explained_variance_ratio_sum": float(np.sum(pca_search.explained_variance_ratio_)),
},
"full": {
"n_samples_fit": 3000,
"n_samples_test": 1000,
"n_components": 32,
"whiten": True,
"random_state": 42,
"explained_variance_ratio_sum": float(np.sum(pca_full.explained_variance_ratio_)),
},
}
all_search_candidates = get_search_candidates()
if method_filter:
unknown_methods = sorted(set(method_filter) - set(all_search_candidates))
if unknown_methods:
raise ValueError(
f"Unknown tuning method(s) {unknown_methods}. "
f"Available: {sorted(all_search_candidates)}"
)
search_seeds = [51, 52, 53] if not quick else [51]
search_particles = 30 if not quick else 5
search_epochs = 80 if not quick else 5
batch_size = 1000 if not quick else 25
search_runs: List[Dict[str, Any]] = list(completed_search_runs.values())
# --- Phase 1: Search ---
if phase in ("all", "search"):
print("\n=== Phase 1: Search (Validation Tuning) ===")
for method, cand_list in all_search_candidates.items():
if method_filter and method not in method_filter:
continue
run_cands = cand_list[:1] if quick else cand_list
for cfg in run_cands:
for seed in search_seeds:
cfg_payload = {
"phase": "search",
"quick": quick,
"device": str(device),
"candidate_label": cfg.candidate_label,
"seed": seed,
"n_particles": search_particles,
"epochs": search_epochs,
"kwargs": cfg.to_optimizer_kwargs(quick=quick),
}
fp_bytes = json.dumps(cfg_payload, sort_keys=True, default=str).encode("utf-8")
cfg_fp = hashlib.sha256(fp_bytes).hexdigest()[:12]
run_id = f"search_{cfg.candidate_label}_seed{seed}_{cfg_fp}"
if run_id in completed_search_runs and not overwrite:
print(f"Skipping completed search run: {run_id}")
continue
print(f"Running {run_id} ({cfg.description})...")
run_res = run_single_experiment(
cfg=cfg,
seed=seed,
x_train=x_inner_tr,
y_train=y_inner_tr,
x_eval=x_inner_val,
y_eval=y_inner_val,
n_particles=search_particles,
epochs=search_epochs,
batch_size=batch_size,
device=device,
quick=quick,
eval_metric_name="val",
data_fp=search_data_fp,
run_type="search",
extra_meta={"run_id": run_id},
)
completed_search_runs[run_id] = run_res
search_runs = list(completed_search_runs.values())
_persist_json_state(
output_json=output_json,
device=device,
quick=quick,
hw_provenance=hw_provenance,
split_fingerprints=split_fingerprints,
pca_provenance=pca_provenance,
all_search_candidates=all_search_candidates,
search_runs=search_runs,
confirmation_runs=list(completed_confirm_runs.values()),
scaling_runs=list(completed_scaling_runs.values()),
)
# Winner Selection Logic (Requires expected completed search seeds per candidate)
required_seeds = 1 if quick else 3
winners: Dict[str, Dict[str, Any]] = {}
candidate_summaries: Dict[str, Dict[str, Any]] = {}
incomplete_methods = []
for method, cand_list in all_search_candidates.items():
cand_scores = []
for cfg in cand_list:
lbl = cfg.candidate_label
matching_runs = [r for r in search_runs if r.get("candidate_label") == lbl and r.get("completed")]
if len(matching_runs) < required_seeds:
continue
val_accs = [r["val_acc"] for r in matching_runs]
val_losses = [r["val_loss"] for r in matching_runs]
mean_acc = float(np.mean(val_accs))
mean_loss = float(np.mean(val_losses))
cand_scores.append({
"candidate_label": lbl,
"cfg": cfg,
"mean_val_loss": mean_loss,
"mean_val_acc": mean_acc,
"n_runs": len(matching_runs),
"stats_acc": calc_stats(val_accs),
"stats_loss": calc_stats(val_losses),
})
candidate_summaries[lbl] = cand_scores[-1]
if cand_scores:
# Rank validation accuracy descending, then validation loss ascending
cand_scores.sort(key=lambda c: (-c["mean_val_acc"], c["mean_val_loss"]))
top = cand_scores[0]
winners[method] = {
"method": method,
"candidate_label": top["candidate_label"],
"description": top["cfg"].description,
"mean_val_loss": top["mean_val_loss"],
"mean_val_acc": top["mean_val_acc"],
"config": top["cfg"].to_optimizer_kwargs(quick=quick),
"cfg": top["cfg"],
}
else:
incomplete_methods.append(method)
if phase in ("all", "confirmation", "scaling") and incomplete_methods:
raise RuntimeError(
f"Cannot proceed to {phase}: Search phase incomplete for method(s): {incomplete_methods}. "
f"Expected {required_seeds} completed search seeds per candidate."
)
if winners:
print("\n--- Search Winners Selected ---")
for m, w in winners.items():
print(f" {m:15s} -> Winner: {w['candidate_label']} (Val Acc: {w['mean_val_acc']*100:.2f}%, Val Loss: {w['mean_val_loss']:.4f})")
# --- Phase 2: Confirmation ---
confirm_seeds = [61, 62, 63, 64, 65] if not quick else [61]
confirm_particles = 30 if not quick else 5
confirm_epochs = 80 if not quick else 5
confirm_runs: List[Dict[str, Any]] = list(completed_confirm_runs.values())
if phase in ("all", "confirmation"):
print("\n=== Phase 2: Confirmation (Held-Out Test Confirmation) ===")
for method, win_info in winners.items():
if method_filter and method not in method_filter:
continue
cfg = win_info["cfg"]
for seed in confirm_seeds:
cfg_payload = {
"phase": "confirmation",
"quick": quick,
"device": str(device),
"candidate_label": cfg.candidate_label,
"seed": seed,
"n_particles": confirm_particles,
"epochs": confirm_epochs,
"kwargs": cfg.to_optimizer_kwargs(quick=quick),
}
fp_bytes = json.dumps(cfg_payload, sort_keys=True, default=str).encode("utf-8")
cfg_fp = hashlib.sha256(fp_bytes).hexdigest()[:12]
run_id = f"confirm_{method}_{cfg.candidate_label}_seed{seed}_{cfg_fp}"
if run_id in completed_confirm_runs and not overwrite:
print(f"Skipping completed confirmation run: {run_id}")
continue
print(f"Running confirmation {run_id} ({method} / {cfg.candidate_label})...")
run_res = run_single_experiment(
cfg=cfg,
seed=seed,
x_train=x_full_tr,
y_train=y_train_3000,
x_eval=x_full_test,
y_eval=y_test_1000,
n_particles=confirm_particles,
epochs=confirm_epochs,
batch_size=batch_size,
device=device,
quick=quick,
eval_metric_name="test",
data_fp=full_data_fp,
run_type="confirmation",
extra_meta={"run_id": run_id},
)
completed_confirm_runs[run_id] = run_res
confirm_runs = list(completed_confirm_runs.values())
_persist_json_state(
output_json=output_json,
device=device,
quick=quick,
hw_provenance=hw_provenance,
split_fingerprints=split_fingerprints,
pca_provenance=pca_provenance,
all_search_candidates=all_search_candidates,
search_runs=search_runs,
confirmation_runs=confirm_runs,
scaling_runs=list(completed_scaling_runs.values()),
)
# --- Phase 3: Particle Scaling ---
scaling_seeds = [71, 72, 73, 74, 75] if not quick else [71]
scaling_runs: List[Dict[str, Any]] = list(completed_scaling_runs.values())
if phase in ("all", "scaling"):
print("\n=== Phase 3: Particle Scaling Study (Adaptive Moment Winner) ===")
am_winner = winners.get("adaptive_moment")
if not am_winner:
raise RuntimeError("Error: No adaptive_moment search winner available for scaling study.")
cfg = am_winner["cfg"]
if quick:
scaling_configs = [
(5, 5, "fixed_epoch"),
(10, 5, "fixed_epoch"),
(5, 5, "fixed_budget"),
(10, 3, "fixed_budget"),
]
else:
scaling_configs = [
(30, 80, "fixed_epoch"),
(60, 80, "fixed_epoch"),
(90, 80, "fixed_epoch"),
(120, 80, "fixed_epoch"),
(30, 80, "fixed_budget"),
(60, 40, "fixed_budget"),
(90, 27, "fixed_budget"),
(120, 20, "fixed_budget"),
]
exec_cache: Dict[Tuple[int, int, int], Dict[str, Any]] = {}
for r in scaling_runs:
key = (r["n_particles"], r["epochs"], r["seed"])
exec_cache[key] = r
for (p_count, ep_count, regimen) in scaling_configs:
for seed in scaling_seeds:
exec_key = (p_count, ep_count, seed)
cfg_payload = {
"phase": "scaling",
"quick": quick,
"device": str(device),
"candidate_label": cfg.candidate_label,
"regimen": regimen,
"seed": seed,
"n_particles": p_count,
"epochs": ep_count,
"kwargs": cfg.to_optimizer_kwargs(quick=quick),
}
fp_bytes = json.dumps(cfg_payload, sort_keys=True, default=str).encode("utf-8")
cfg_fp = hashlib.sha256(fp_bytes).hexdigest()[:12]
run_id = f"scaling_{p_count}p_{ep_count}e_{regimen}_seed{seed}_{cfg_fp}"
if run_id in completed_scaling_runs and not overwrite:
print(f"Skipping completed scaling run: {run_id}")
continue
if exec_key in exec_cache:
print(f"Reusing deduplicated run for {run_id} ({p_count} particles, {ep_count} epochs, seed {seed})...")
existing_res = dict(exec_cache[exec_key])
existing_res["run_id"] = run_id
existing_res["regimen"] = regimen
run_res = existing_res
else:
print(f"Running scaling {run_id} ({p_count} particles, {ep_count} epochs, {regimen}, seed {seed})...")
run_res = run_single_experiment(
cfg=cfg,
seed=seed,
x_train=x_full_tr,
y_train=y_train_3000,
x_eval=x_full_test,
y_eval=y_test_1000,
n_particles=p_count,
epochs=ep_count,
batch_size=batch_size,
device=device,
quick=quick,
eval_metric_name="test",
data_fp=full_data_fp,
run_type="scaling",
extra_meta={"run_id": run_id, "regimen": regimen},
)
exec_cache[exec_key] = run_res
completed_scaling_runs[run_id] = run_res
scaling_runs = list(completed_scaling_runs.values())
_persist_json_state(
output_json=output_json,
device=device,
quick=quick,
hw_provenance=hw_provenance,
split_fingerprints=split_fingerprints,
pca_provenance=pca_provenance,
all_search_candidates=all_search_candidates,
search_runs=search_runs,
confirmation_runs=confirm_runs,
scaling_runs=scaling_runs,
)
write_tuning_csvs(search_runs, confirm_runs, scaling_runs, result_dir)
if phase in ("all", "plots"):
render_tuning_plots(search_runs, confirm_runs, scaling_runs, winners, figure_dir)
print("\nMNIST Tuning Study Complete!")
print(f"- Primary JSON: {output_json}")
print(f"- Search CSV: {result_dir / 'pso_v4_tuning_search.csv'}")
print(f"- Confirmation CSV: {result_dir / 'pso_v4_tuning_confirmation.csv'}")
print(f"- Scaling CSV: {result_dir / 'pso_v4_particle_scaling.csv'}")
print(f"- Figures in: {figure_dir}")
def main():
parser = argparse.ArgumentParser(description="MNIST PSO Tuning & Particle Scaling Study Suite")
parser.add_argument(
"--phase",
choices=["all", "search", "confirmation", "scaling", "plots"],
default="all",
help="Study phase to execute (default: all)",
)
parser.add_argument(
"--device",
type=str,
default=None,
help="Execution device ('mps', 'cuda', 'cpu')",
)
parser.add_argument(
"--quick",
action="store_true",
help="Run reduced quick smoke test across all phases",
)
parser.add_argument(
"--overwrite",
action="store_true",
help="Overwrite existing completed run checkpoints and JSON results",
)
parser.add_argument(
"--methods",
type=str,
default=None,
help="Comma-separated method filter for search/confirmation reruns",
)
parser.add_argument(
"--output-json",
type=Path,
default=Path("benchmark_results/pso_v4_tuning.json"),
help="Path to output JSON result file",
)
parser.add_argument(
"--result-dir",
type=Path,
default=Path("benchmark_results"),
help="Directory to save CSV report artifacts",
)
parser.add_argument(
"--figure-dir",
type=Path,
default=Path("history_plt"),
help="Directory to save PNG figure artifacts",
)
args = parser.parse_args()
method_filter = (
[method.strip() for method in args.methods.split(",") if method.strip()]
if args.methods
else None
)
run_tuning_study(
phase=args.phase,
device_name=args.device,
quick=args.quick,
method_filter=method_filter,
overwrite=args.overwrite,
output_json=args.output_json,
result_dir=args.result_dir,
figure_dir=args.figure_dir,
)
if __name__ == "__main__":
main()
Reference in New Issue View Git Blame Copy Permalink