408 lines
17 KiB
Python
408 lines
17 KiB
Python
#!/usr/bin/env python
|
|
"""Baseline PROTEAN smoke/run entry point.
|
|
|
|
Runs the frozen baseline pipeline end to end for one ``(dataset, alpha, seed)``:
|
|
load -> clean -> split -> preprocess -> Dirichlet partition -> federated training
|
|
-> local prototypes -> server aggregation -> global redistribution -> evaluation
|
|
-> structured artifacts.
|
|
|
|
This runner defaults to a small configuration suitable for a smoke test. A full
|
|
scientific sweep must be launched separately once the baseline is verified.
|
|
|
|
Usage (smoke)::
|
|
|
|
python scripts/run_baseline.py --dataset xiiotid --alpha 0.5 --seed 0 \
|
|
--max-rows 4000 --clients 4 --rounds 2 --local-epochs 1
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(REPO_ROOT / "src"))
|
|
|
|
from protean_research.data import manifest as manifest_mod # noqa: E402
|
|
from protean_research.data import partition as partition_mod # noqa: E402
|
|
from protean_research.data import rare_unseen # noqa: E402
|
|
from protean_research.aggregation import validate_beta # noqa: E402
|
|
from protean_research.data.characterization import characterize_support # noqa: E402
|
|
from protean_research.data.pipeline import ( # noqa: E402
|
|
create_partition,
|
|
prepare_dataset,
|
|
preprocessing_content_hash,
|
|
)
|
|
from protean_research.evaluation.metrics import ( # noqa: E402
|
|
classification_metrics,
|
|
rare_class_metrics,
|
|
)
|
|
from protean_research.federated import TrainingConfig, run_federated # noqa: E402
|
|
from protean_research.models.protean_cnn import extract_embeddings # noqa: E402
|
|
from protean_research.models.prototypes import nearest_prototype_predict # noqa: E402
|
|
from protean_research.utils.config import ( # noqa: E402
|
|
expand_path,
|
|
load_datasets_config,
|
|
load_yaml,
|
|
)
|
|
|
|
|
|
def parse_args(argv=None) -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--dataset", required=True)
|
|
parser.add_argument("--alpha", type=float, required=True)
|
|
parser.add_argument("--seed", type=int, default=0)
|
|
parser.add_argument("--clients", type=int, default=10)
|
|
parser.add_argument("--rounds", type=int, default=10)
|
|
parser.add_argument("--local-epochs", type=int, default=3)
|
|
parser.add_argument("--batch-size", type=int, default=256)
|
|
parser.add_argument("--lr", type=float, default=1e-3)
|
|
parser.add_argument("--mu", type=float, default=0.1)
|
|
parser.add_argument("--lam", type=float, default=1.0)
|
|
parser.add_argument(
|
|
"--aggregation-mode",
|
|
choices=["original_protean", "support_weighted"],
|
|
default="original_protean",
|
|
help="server-side prototype aggregation mode (default: frozen baseline).",
|
|
)
|
|
parser.add_argument(
|
|
"--beta",
|
|
type=float,
|
|
default=None,
|
|
help="support-weighting exponent in [0, 1]; required for support_weighted.",
|
|
)
|
|
parser.add_argument(
|
|
"--max-rows",
|
|
type=int,
|
|
default=None,
|
|
help="SMOKE ONLY: limit how many raw rows are read; never use for science.",
|
|
)
|
|
parser.add_argument("--sample-from", type=int, default=200000,
|
|
help="SMOKE ONLY: rows to read before random subsampling.")
|
|
parser.add_argument("--partitions-dir", default=None)
|
|
parser.add_argument("--processed-dir", default=None)
|
|
parser.add_argument("--results-dir", default=None)
|
|
parser.add_argument("--analysis-dir", default=None)
|
|
return parser.parse_args(argv)
|
|
|
|
|
|
def write_json(path: Path, payload) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with path.open("w", encoding="utf-8") as handle:
|
|
json.dump(payload, handle, indent=2, sort_keys=True)
|
|
handle.write("\n")
|
|
|
|
|
|
def write_csv(path: Path, rows: list[dict]) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
if not rows:
|
|
path.write_text("", encoding="utf-8")
|
|
return
|
|
fieldnames = list(rows[0].keys())
|
|
with path.open("w", encoding="utf-8", newline="") as handle:
|
|
writer = csv.DictWriter(handle, fieldnames=fieldnames)
|
|
writer.writeheader()
|
|
writer.writerows(rows)
|
|
|
|
|
|
def main(argv=None) -> int:
|
|
args = parse_args(argv)
|
|
|
|
if args.aggregation_mode == "support_weighted":
|
|
if args.beta is None:
|
|
raise SystemExit(
|
|
"--beta is required when --aggregation-mode support_weighted"
|
|
)
|
|
beta = validate_beta(args.beta)
|
|
else:
|
|
if args.beta is not None:
|
|
raise SystemExit(
|
|
"--beta is only valid with --aggregation-mode support_weighted"
|
|
)
|
|
beta = None
|
|
|
|
datasets_cfg = load_datasets_config()
|
|
protocol = load_yaml(REPO_ROOT / "configs" / "experiments" / "experiment_protocol.yaml")
|
|
|
|
artifacts_cfg = protocol.get("artifacts", {})
|
|
partitions_dir = expand_path(args.partitions_dir or artifacts_cfg.get("partition_manifest_dir", "artifacts/partitions"))
|
|
processed_dir = expand_path(args.processed_dir or artifacts_cfg.get("processed_dir", "artifacts/processed"))
|
|
analysis_dir = expand_path(args.analysis_dir or artifacts_cfg.get("analysis_dir", "artifacts/analysis"))
|
|
results_root = expand_path(args.results_dir or artifacts_cfg.get("results_dir", "results"))
|
|
|
|
# Smoke runs must never pollute the tracked partition/analysis directories.
|
|
if args.max_rows is not None:
|
|
if args.partitions_dir is None:
|
|
partitions_dir = processed_dir / "smoke" / "partitions"
|
|
if args.analysis_dir is None:
|
|
analysis_dir = processed_dir / "smoke" / "analysis"
|
|
if args.results_dir is None:
|
|
results_root = results_root / "smoke"
|
|
|
|
if args.max_rows is not None:
|
|
from protean_research.data import loading as loading_mod
|
|
import numpy as np
|
|
|
|
spec = loading_mod.load_spec(args.dataset, datasets_cfg)
|
|
pool = loading_mod.read_raw(spec, nrows=args.sample_from)
|
|
rng = np.random.default_rng(args.seed)
|
|
take = min(args.max_rows, len(pool))
|
|
chosen = np.sort(rng.choice(len(pool), size=take, replace=False))
|
|
pool = pool.iloc[chosen].reset_index(drop=True)
|
|
bundle = prepare_dataset(
|
|
args.dataset, seed=args.seed, datasets_cfg=datasets_cfg, raw_frame=pool
|
|
)
|
|
else:
|
|
bundle = prepare_dataset(args.dataset, seed=args.seed, datasets_cfg=datasets_cfg)
|
|
|
|
# ---- partition (create or reload) -------------------------------------
|
|
pre_hash = preprocessing_content_hash(bundle.preprocessor)
|
|
manifest_file = manifest_mod.manifest_path(partitions_dir, args.dataset, args.alpha, args.seed)
|
|
if manifest_file.exists():
|
|
manifest = manifest_mod.load_manifest(manifest_file)
|
|
split, assignments, support = manifest_mod.reproduce_train_assignments(manifest, bundle.labels)
|
|
# Reuse the manifest's split for pairing; verify it matches the bundle.
|
|
bundle.split = split
|
|
else:
|
|
manifest, assignments, support = create_partition(
|
|
bundle, alpha=args.alpha, seed=args.seed, num_clients=args.clients,
|
|
partition_dir=partitions_dir, preprocessing_hash=pre_hash,
|
|
)
|
|
manifest_reload_ok = True
|
|
try:
|
|
manifest_mod.reproduce_train_assignments(manifest, bundle.labels)
|
|
except Exception: # pragma: no cover - surfaced in metadata
|
|
manifest_reload_ok = False
|
|
|
|
labels = bundle.labels
|
|
train_features = bundle.features[bundle.split.train]
|
|
train_labels = labels[bundle.split.train]
|
|
val_features = bundle.features[bundle.split.validation]
|
|
val_labels = labels[bundle.split.validation]
|
|
test_features = bundle.features[bundle.split.test]
|
|
test_labels = labels[bundle.split.test]
|
|
|
|
benign_index = bundle.cleaned.benign_index
|
|
per_client_rare = rare_unseen.rare_class_ids(support, benign_index)
|
|
union_rare = sorted({c for row in per_client_rare for c in row})
|
|
|
|
cfg = TrainingConfig(
|
|
rounds=args.rounds,
|
|
local_epochs=args.local_epochs,
|
|
batch_size=args.batch_size,
|
|
learning_rate=args.lr,
|
|
mu=args.mu,
|
|
lam=args.lam,
|
|
num_clients=args.clients,
|
|
aggregation_mode=args.aggregation_mode,
|
|
beta=beta,
|
|
)
|
|
result = run_federated(
|
|
train_features=train_features,
|
|
train_labels=train_labels,
|
|
assignments=assignments,
|
|
val_features=val_features,
|
|
val_labels=val_labels,
|
|
num_classes=bundle.num_classes,
|
|
input_dim=bundle.features.shape[1],
|
|
config=cfg,
|
|
init_seed_parts=(args.dataset, f"seed={args.seed}"),
|
|
rare_class_ids=union_rare,
|
|
)
|
|
|
|
model = result["model"]
|
|
global_prototypes = result["global_prototypes"]
|
|
|
|
# ---- global test evaluation ------------------------------------------
|
|
test_embeddings = extract_embeddings(model, test_features)
|
|
test_pred = nearest_prototype_predict(test_embeddings, global_prototypes)
|
|
test_metrics = classification_metrics(test_labels, test_pred, bundle.num_classes)
|
|
test_rare = rare_class_metrics(test_labels, test_pred, bundle.num_classes, union_rare)
|
|
|
|
# ---- per-client evaluation -------------------------------------------
|
|
eval_assignments, eval_rng = partition_mod.partition_eval_split(
|
|
test_labels, args.clients, args.alpha, args.dataset, args.seed, "test"
|
|
)
|
|
per_client_rows = []
|
|
for client_id in range(args.clients):
|
|
idx = eval_assignments[client_id]
|
|
if len(idx) == 0:
|
|
per_client_rows.append({
|
|
"client_id": client_id, "n_samples": 0, "accuracy": None,
|
|
"macro_f1": None, "rare_class_macro_f1": None, "rare_class_recall": None,
|
|
})
|
|
continue
|
|
emb = extract_embeddings(model, test_features[idx])
|
|
pred = nearest_prototype_predict(emb, global_prototypes)
|
|
metrics = classification_metrics(test_labels[idx], pred, bundle.num_classes)
|
|
rare_ids = per_client_rare[client_id]
|
|
if rare_ids:
|
|
rare = rare_class_metrics(test_labels[idx], pred, bundle.num_classes, rare_ids)
|
|
else:
|
|
rare = {"rare_class_macro_f1": None, "rare_class_recall": None}
|
|
per_client_rows.append({
|
|
"client_id": client_id,
|
|
"n_samples": int(len(idx)),
|
|
"accuracy": metrics["accuracy"],
|
|
"macro_f1": metrics["macro_f1"],
|
|
"rare_class_macro_f1": rare["rare_class_macro_f1"],
|
|
"rare_class_recall": rare["rare_class_recall"],
|
|
})
|
|
|
|
# ---- global validation evaluation (SELECTION ONLY) -------------------
|
|
# This block never influences training and is never used to tune any
|
|
# hyperparameter except the prototype-aggregation beta, which is selected
|
|
# on validation metrics only. The test split is untouched here.
|
|
val_embeddings = extract_embeddings(model, val_features)
|
|
val_pred = nearest_prototype_predict(val_embeddings, global_prototypes)
|
|
val_metrics = classification_metrics(val_labels, val_pred, bundle.num_classes)
|
|
val_rare = rare_class_metrics(val_labels, val_pred, bundle.num_classes, union_rare)
|
|
|
|
val_eval_assignments, _ = partition_mod.partition_eval_split(
|
|
val_labels, args.clients, args.alpha, args.dataset, args.seed, "validation"
|
|
)
|
|
val_per_client_rows = []
|
|
for client_id in range(args.clients):
|
|
idx = val_eval_assignments[client_id]
|
|
if len(idx) == 0:
|
|
val_per_client_rows.append({
|
|
"client_id": client_id, "n_samples": 0, "accuracy": None,
|
|
"macro_f1": None, "rare_class_macro_f1": None, "rare_class_recall": None,
|
|
})
|
|
continue
|
|
emb = extract_embeddings(model, val_features[idx])
|
|
pred = nearest_prototype_predict(emb, global_prototypes)
|
|
metrics = classification_metrics(val_labels[idx], pred, bundle.num_classes)
|
|
rare_ids = per_client_rare[client_id]
|
|
if rare_ids:
|
|
rare = rare_class_metrics(val_labels[idx], pred, bundle.num_classes, rare_ids)
|
|
else:
|
|
rare = {"rare_class_macro_f1": None, "rare_class_recall": None}
|
|
val_per_client_rows.append({
|
|
"client_id": client_id,
|
|
"n_samples": int(len(idx)),
|
|
"accuracy": metrics["accuracy"],
|
|
"macro_f1": metrics["macro_f1"],
|
|
"rare_class_macro_f1": rare["rare_class_macro_f1"],
|
|
"rare_class_recall": rare["rare_class_recall"],
|
|
})
|
|
|
|
def _mean_over_present(rows, key):
|
|
values = [float(row[key]) for row in rows if row.get(key) is not None]
|
|
return float(sum(values) / len(values)) if values else None
|
|
|
|
val_selection = {
|
|
"val_rare_class_macro_f1_global": val_rare["rare_class_macro_f1"],
|
|
"val_rare_class_recall_global": val_rare["rare_class_recall"],
|
|
"val_macro_f1_global": val_metrics["macro_f1"],
|
|
"val_weighted_f1_global": val_metrics["weighted_f1"],
|
|
"val_accuracy_global": val_metrics["accuracy"],
|
|
"val_rare_class_macro_f1_per_client_mean": _mean_over_present(
|
|
val_per_client_rows, "rare_class_macro_f1"
|
|
),
|
|
"val_rare_class_recall_per_client_mean": _mean_over_present(
|
|
val_per_client_rows, "rare_class_recall"
|
|
),
|
|
"val_n_eval_clients": sum(1 for row in val_per_client_rows if row["n_samples"] > 0),
|
|
"selection_metric": "rare_class_macro_f1_per_client_mean",
|
|
"selection_split": "validation",
|
|
"test_used_for_selection": False,
|
|
}
|
|
|
|
# ---- structured artifacts --------------------------------------------
|
|
mode_tag = args.aggregation_mode
|
|
if args.aggregation_mode == "support_weighted":
|
|
mode_tag = f"support_weighted__beta{beta:g}"
|
|
run_name = f"{args.dataset}__alpha{args.alpha:g}__seed{args.seed}__{mode_tag}"
|
|
run_dir = results_root / run_name
|
|
run_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
metadata = {
|
|
"run_name": run_name,
|
|
"dataset": args.dataset,
|
|
"alpha": args.alpha,
|
|
"seed": args.seed,
|
|
"num_clients": args.clients,
|
|
"num_classes": bundle.num_classes,
|
|
"class_names": bundle.class_names,
|
|
"benign_class": bundle.spec.benign_label,
|
|
"aggregation_mode": args.aggregation_mode,
|
|
"beta": beta,
|
|
"training": {
|
|
"rounds": args.rounds,
|
|
"local_epochs": args.local_epochs,
|
|
"batch_size": args.batch_size,
|
|
"learning_rate": args.lr,
|
|
"mu": args.mu,
|
|
"lambda": args.lam,
|
|
},
|
|
"smoke": args.max_rows is not None,
|
|
"max_rows": args.max_rows,
|
|
"manifest_path": str(manifest_file),
|
|
"manifest_content_hash": manifest_mod.manifest_content_hash(manifest),
|
|
"preprocessing_hash": pre_hash,
|
|
"manifest_reload_ok": manifest_reload_ok,
|
|
"repo_root": str(REPO_ROOT),
|
|
}
|
|
write_json(run_dir / "metadata.json", metadata)
|
|
write_json(run_dir / "history.json", result["history"])
|
|
write_json(run_dir / "test_global_metrics.json", {
|
|
"accuracy": test_metrics["accuracy"],
|
|
"macro_f1": test_metrics["macro_f1"],
|
|
"weighted_f1": test_metrics["weighted_f1"],
|
|
**test_rare,
|
|
})
|
|
write_json(run_dir / "test_confusion_matrix.json", test_metrics["confusion_matrix"])
|
|
write_csv(run_dir / "test_per_class.csv", [
|
|
{**row, "class_name": bundle.class_names[row["class_id"]]}
|
|
for row in test_metrics["per_class"]
|
|
])
|
|
write_csv(run_dir / "per_client.csv", per_client_rows)
|
|
write_json(run_dir / "val_global_metrics.json", {
|
|
"accuracy": val_metrics["accuracy"],
|
|
"macro_f1": val_metrics["macro_f1"],
|
|
"weighted_f1": val_metrics["weighted_f1"],
|
|
**val_rare,
|
|
})
|
|
write_json(run_dir / "val_confusion_matrix.json", val_metrics["confusion_matrix"])
|
|
write_csv(run_dir / "val_per_class.csv", [
|
|
{**row, "class_name": bundle.class_names[row["class_id"]]}
|
|
for row in val_metrics["per_class"]
|
|
])
|
|
write_csv(run_dir / "val_per_client.csv", val_per_client_rows)
|
|
write_json(run_dir / "val_selection.json", val_selection)
|
|
write_csv(run_dir / "train_support_Nij.csv", [
|
|
{"client_id": i, **{bundle.class_names[c]: int(support[i, c]) for c in range(bundle.num_classes)}}
|
|
for i in range(args.clients)
|
|
])
|
|
write_csv(run_dir / "characterization.csv", characterize_support(support, bundle.class_names))
|
|
write_csv(run_dir / "rare_unseen.csv", rare_unseen.rare_unseen_rows(
|
|
support, bundle.class_names, benign_index, args.dataset, args.alpha, args.seed
|
|
))
|
|
|
|
processed_dir.mkdir(parents=True, exist_ok=True)
|
|
bundle.preprocessor.save(processed_dir / f"{args.dataset}__seed{args.seed}__preprocessor.json")
|
|
|
|
snapshot = result["snapshot"]
|
|
if snapshot is not None:
|
|
write_json(analysis_dir / f"{run_name}__prototype_snapshot.json", snapshot)
|
|
|
|
summary = {
|
|
"run": run_name,
|
|
"test_accuracy": test_metrics["accuracy"],
|
|
"test_macro_f1": test_metrics["macro_f1"],
|
|
"test_rare_class_macro_f1": test_rare["rare_class_macro_f1"],
|
|
"n_global_prototypes": len(global_prototypes),
|
|
"history_last": result["history"][-1] if result["history"] else None,
|
|
"results_dir": str(run_dir),
|
|
}
|
|
print(json.dumps(summary, indent=2, sort_keys=True))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|