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/generate_partitions.py
+124
View File
@@ -0,0 +1,124 @@
#!/usr/bin/env python
"""Generate deterministic Dirichlet partition manifests for all conditions.
This does NOT train anything. It loads each dataset once (labels only), computes
the frozen stratified split, generates the Dirichlet partition for every
``(alpha, seed)`` and persists a manifest plus characterization/rare-unseen
metadata. Training later loads these manifests; it never regenerates a partition.
Usage::
python scripts/generate_partitions.py --datasets xiiotid 5gnidd \
--alphas 0.75 0.5 0.25 --seeds 0 1 2 --clients 10
"""
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 loading, manifest as manifest_mod, partition as partition_mod # noqa: E402
from protean_research.data import rare_unseen, splits as splits_mod # noqa: E402
from protean_research.data.characterization import characterize_support # noqa: E402
from protean_research.data.pipeline import ( # noqa: E402
TEST_FRACTION,
TRAIN_FRACTION,
VALIDATION_FRACTION,
)
from protean_research.utils.config import expand_path, load_datasets_config, load_yaml # noqa: E402
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
with path.open("w", encoding="utf-8", newline="") as handle:
writer = csv.DictWriter(handle, fieldnames=list(rows[0].keys()))
writer.writeheader()
writer.writerows(rows)
def main(argv=None) -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--datasets", nargs="+", default=["xiiotid", "5gnidd"])
parser.add_argument("--alphas", nargs="+", type=float, default=[0.75, 0.5, 0.25])
parser.add_argument("--seeds", nargs="+", type=int, default=[0, 1, 2])
parser.add_argument("--clients", type=int, default=10)
parser.add_argument("--partitions-dir", default=None)
parser.add_argument("--analysis-dir", default=None)
args = parser.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"))
analysis_dir = expand_path(args.analysis_dir or artifacts_cfg.get("analysis_dir", "artifacts/analysis"))
created = []
for dataset_key in args.datasets:
spec = loading.load_spec(dataset_key, datasets_cfg)
cleaned = loading.clean_frame(loading.read_raw(spec), spec)
for alpha in args.alphas:
for seed in args.seeds:
split = splits_mod.stratified_split(
cleaned.labels, dataset_key, seed,
train_fraction=TRAIN_FRACTION,
validation_fraction=VALIDATION_FRACTION,
test_fraction=TEST_FRACTION,
)
labels_train = cleaned.labels[split.train]
class_indices = partition_mod.class_index_map(labels_train, len(cleaned.class_names))
rng_seed = partition_mod.partition_rng_seed(dataset_key, alpha, seed)
assignments = partition_mod.dirichlet_partition(
class_indices, args.clients, alpha, rng_seed
)
manifest = manifest_mod.build_manifest(
dataset_key=dataset_key,
dataset_name=spec.name,
alpha=alpha,
seed=seed,
num_clients=args.clients,
class_names=cleaned.class_names,
benign_class=spec.benign_label,
labels=cleaned.labels,
split=split,
assignments=assignments,
rng_seed=rng_seed,
train_fraction=TRAIN_FRACTION,
validation_fraction=VALIDATION_FRACTION,
test_fraction=TEST_FRACTION,
)
path = manifest_mod.manifest_path(partitions_dir, dataset_key, alpha, seed)
manifest_mod.save_manifest(path, manifest)
# Verify exact reload before declaring success.
manifest_mod.reproduce_train_assignments(manifest, cleaned.labels)
support = manifest["support"]
import numpy as np
support_arr = np.asarray(support, dtype=np.int64)
tag = f"{dataset_key}__alpha{alpha:g}__seed{seed}"
write_csv(analysis_dir / f"{tag}__characterization.csv",
characterize_support(support_arr, cleaned.class_names))
write_csv(analysis_dir / f"{tag}__rare_unseen.csv",
rare_unseen.rare_unseen_rows(
support_arr, cleaned.class_names, cleaned.benign_index,
dataset_key, alpha, seed))
created.append({"dataset": dataset_key, "alpha": alpha, "seed": seed,
"path": str(path), "partition_hash": manifest["partition_hash"]})
print(f"created {path.name}")
print(json.dumps({"created": len(created)}, indent=2))
return 0
if __name__ == "__main__":
raise SystemExit(main())
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())