diff options
| -rw-r--r-- | experiments/plot_main_figures.py | 104 |
1 files changed, 97 insertions, 7 deletions
diff --git a/experiments/plot_main_figures.py b/experiments/plot_main_figures.py index 9306d90..e171181 100644 --- a/experiments/plot_main_figures.py +++ b/experiments/plot_main_figures.py @@ -28,6 +28,7 @@ COLORS = { } MARKERS = {"BP": "*", "FA": "s", "DFA": "^", "SDIL": "o", "EP": "D"} DEPTHS = [5, 10, 20, 30, 60] +EXPECTED_SEEDS = set(range(5)) def read_many(patterns): @@ -58,17 +59,96 @@ def require_clean(rows): raise RuntimeError("dirty or unknown provenance:\n" + "\n".join(bad)) -def require_cells(groups, expected, n, label, strict): +def require_cells(groups, expected, label, strict): missing = [] for key in expected: - count = len(groups.get(key, [])) - if count != n: - missing.append(f"{key}: expected {n}, found {count}") + rows = groups.get(key, []) + seeds = [r["args"].get("seed") for r in rows] + if len(seeds) != len(set(seeds)): + missing.append(f"{key}: duplicate seeds {sorted(seeds)}") + elif set(seeds) != EXPECTED_SEEDS: + missing.append( + f"{key}: expected seeds {sorted(EXPECTED_SEEDS)}, found {sorted(seeds)}") if missing and strict: raise RuntimeError(f"incomplete {label}:\n" + "\n".join(missing)) return missing +def require_value(errors, row, key, actual, expected): + if actual != expected: + errors.append(f"{row['_path']}: {key} expected {expected!r}, found {actual!r}") + + +def require_valid_finals(rows): + errors = [] + for row in rows: + final = row.get("final", {}) + acc, wall = final.get("test_acc"), final.get("wall_s") + if not isinstance(acc, (int, float)) or not math.isfinite(acc) or not 0 <= acc <= 1: + errors.append(f"{row['_path']}: invalid final test_acc {acc!r}") + if not isinstance(wall, (int, float)) or not math.isfinite(wall) or wall <= 0: + errors.append(f"{row['_path']}: invalid final wall_s {wall!r}") + if errors: + raise RuntimeError("invalid result values:\n" + "\n".join(errors)) + + +def require_scaling_protocol(rows): + errors = [] + for row in rows: + args = row["args"] + for key, expected in (("dataset", "cifar10"), ("width", 64), + ("epochs", 5), ("batch_size", 128), + ("act", "tanh"), ("residual", 1), + ("device", "cuda")): + require_value(errors, row, key, args.get(key), expected) + name = method(row) + require_value(errors, row, "depth", args.get("depth"), + args.get("depth") if args.get("depth") in DEPTHS else "one of " + repr(DEPTHS)) + require_value(errors, row, "eta", args.get("eta"), 0.02 if name == "SDIL" else 0.05) + if name == "SDIL": + require_value(errors, row, "pert_mode", args.get("pert_mode"), "simultaneous") + require_value(errors, row, "pert_ndirs", args.get("pert_ndirs"), 16) + if errors: + raise RuntimeError("mixed CIFAR scaling protocols:\n" + "\n".join(errors)) + + +def require_ep_protocol(groups): + errors = [] + dynamics = { + 1: {"epochs": 25, "beta": 0.5, "free_steps": 20, + "nudge_steps": 4, "learning_rates": [0.1, 0.05]}, + 2: {"epochs": 60, "beta": 1.0, "free_steps": 150, + "nudge_steps": 6, "learning_rates": [0.4, 0.1, 0.01]}, + } + for (name, depth), rows in groups.items(): + for row in rows: + args = row["args"] + for key, expected in (("dataset", "mnist"), ("width", 500), + ("depth", depth), ("train_examples", 50000), + ("device", "cuda"), + ("epochs", dynamics[depth]["epochs"])): + require_value(errors, row, key, args.get(key), expected) + if name == "EP": + require_value(errors, row, "batch_size", args.get("batch_size"), 20) + require_value(errors, row, "ep_persistent", args.get("ep_persistent"), 1) + resolved = row.get("resolved_protocol", {}) + for key in ("beta", "free_steps", "nudge_steps", "learning_rates"): + require_value(errors, row, f"resolved_protocol.{key}", + resolved.get(key), dynamics[depth][key]) + require_value(errors, row, "persistent_particles", + resolved.get("persistent_particles"), True) + else: + require_value(errors, row, "batch_size", args.get("batch_size"), 128) + require_value(errors, row, "act", args.get("act"), "tanh") + require_value(errors, row, "residual", args.get("residual"), 0) + require_value(errors, row, "eta", args.get("eta"), 0.05) + if name == "SDIL": + require_value(errors, row, "pert_mode", args.get("pert_mode"), "simultaneous") + require_value(errors, row, "pert_ndirs", args.get("pert_ndirs"), 16) + if errors: + raise RuntimeError("mixed exact-EP protocols:\n" + "\n".join(errors)) + + def scaling_records(root): return read_many([ os.path.join(root, "scale_v3_cifar10_*.json"), @@ -141,13 +221,16 @@ def plot_tradeoff(root, outdir, strict): scale = scaling_records(root) ep_rows, ep = ep_groups(root) require_clean(scale + ep_rows) + require_valid_finals(scale + ep_rows) sg = scaling_groups(scale) + require_scaling_protocol(scale) + require_ep_protocol(ep) miss_scale = require_cells( sg, [(m, d) for m in ("BP", "FA", "DFA", "SDIL") for d in DEPTHS], - 5, "CIFAR scaling", strict) + "CIFAR scaling", strict) miss_ep = require_cells( ep, [(m, d) for m in ("BP", "DFA", "SDIL", "EP") for d in (1, 2)], - 5, "exact EP comparison", strict) + "exact EP comparison", strict) fig, axes = plt.subplots(1, 2, figsize=(7.35, 3.15), gridspec_kw={"wspace": 0.30}) @@ -228,10 +311,12 @@ def early_alignment(row): def plot_scaling(root, outdir, strict): rows = scaling_records(root) require_clean(rows) + require_valid_finals(rows) + require_scaling_protocol(rows) groups = scaling_groups(rows) missing = require_cells( groups, [(m, d) for m in ("BP", "FA", "DFA", "SDIL") for d in DEPTHS], - 5, "CIFAR scaling", strict) + "CIFAR scaling", strict) fig, axes = plt.subplots(1, 3, figsize=(7.35, 2.75), gridspec_kw={"wspace": 0.38}) @@ -323,6 +408,11 @@ def main(): source_paths = sorted({r["_path"] for r in rows1 + rows2}) manifest = { "strict": strict, + "required_seeds": sorted(EXPECTED_SEEDS), + "protocols": { + "cifar_scaling": "width64 residual tanh; 5 epochs; d5/10/20/30/60", + "ep_comparison": "exact d1/d2 width500; canonical EP dynamics", + }, "missing_cells": missing1 + missing2, "sources": [{"path": p, "sha256": file_hash(p)} for p in source_paths], "outputs": ["figure1_pareto.pdf", "figure1_pareto.png", |