willtanoe/thesis-prototype-agregation
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: implement reproducible PROTEAN baseline pipeline

Browse Source
This commit is contained in:
willtanoe committed 2026-09-24 17:39:00 +00:00
1 parent 2cc6a9859e
commit 0aa8e2ba79
93 files changed
+8173 -658
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

scripts/run_baseline.py
+301
View File
@@ -0,0 +1,301 @@
#!/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.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(
"--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)
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,
)
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"],
})
# ---- structured artifacts --------------------------------------------
run_name = f"{args.dataset}__alpha{args.alpha:g}__seed{args.seed}__original_protean"
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": "original_protean",
"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_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())