1 files changed, 83 insertions, 0 deletions
diff --git a/experiments/zoo.py b/experiments/zoo.py
new file mode 100644
index 0000000..8462be9
--- /dev/null
+++ b/experiments/zoo.py
@@ -0,0 +1,83 @@
+"""Train the local-learning baseline 'zoo' and report test accuracy, for the
+comparison table alongside SDIL/DFA/BP. FA and EP are the working ones; PEPITA
+and FF are included but currently undertuned.
+Usage: python experiments/zoo.py --dataset mnist --methods fa,ep --epochs 10
+"""
+import argparse
+import json
+import os
+import sys
+import time
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from sdil.local_baselines import FANet, PEPITANet, FFNet, EPNet
+from sdil.baselines import evaluate
+from sdil.data import get_dataset, onehot
+
+
+def run_method(m, args, device):
+    tr, te, n_in, n_out = get_dataset(args.dataset, args.batch_size, device=device)
+    sizes = [n_in] + [args.width] * args.depth + [10]
+    t0 = time.time()
+    if m == "fa":
+        net = FANet(sizes, act="tanh", device=device, seed=args.seed)
+        for ep in range(args.epochs):
+            for x, y in tr:
+                net.fa_step(x, y, onehot(y, 10, device=device), args.eta, 0.9)
+        acc = evaluate(net, te)[0]
+    elif m == "pepita":
+        net = PEPITANet(sizes, act="tanh", device=device, seed=args.seed, f_scale=0.5)
+        for ep in range(args.epochs):
+            for x, y in tr:
+                net.pepita_step(x, y, onehot(y, 10, device=device), args.eta, 0.9)
+        acc = evaluate(net, te)[0]
+    elif m == "ff":
+        net = FFNet(sizes, act="relu", device=device, seed=args.seed, threshold=2.0, overlay_val=10.0)
+        for ep in range(args.epochs):
+            for x, y in tr:
+                net.train_step(x, y, args.eta)
+        acc = net.evaluate(te)[0]
+    elif m == "ep":
+        net = EPNet([n_in, args.width, 10], device=device, seed=args.seed,
+                    beta=0.5, dt=0.5, T_free=20, T_nudge=8)
+        for ep in range(args.epochs):
+            for x, y in tr:
+                net.train_step(x, y, onehot(y, 10, device=device), args.eta)
+        acc = net.evaluate(te)[0]
+    else:
+        raise ValueError(m)
+    return acc, time.time() - t0
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--dataset", default="mnist")
+    p.add_argument("--methods", default="fa,ep,pepita,ff")
+    p.add_argument("--depth", type=int, default=2)
+    p.add_argument("--width", type=int, default=500)
+    p.add_argument("--epochs", type=int, default=10)
+    p.add_argument("--batch_size", type=int, default=64)
+    p.add_argument("--eta", type=float, default=0.1)
+    p.add_argument("--seed", type=int, default=0)
+    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    p.add_argument("--outdir", default="results")
+    p.add_argument("--tag", default="zoo")
+    args = p.parse_args()
+
+    # sensible per-method default LRs if the shared one is off
+    lrs = {"fa": 0.05, "ep": 0.1, "pepita": 0.05, "ff": 0.03}
+    out = {"args": vars(args), "acc": {}}
+    for m in args.methods.split(","):
+        args.eta = lrs.get(m, args.eta)
+        acc, secs = run_method(m, args, args.device)
+        out["acc"][m] = acc
+        print(f"[{args.dataset}] {m}: test_acc {acc:.4f} ({secs:.0f}s)", flush=True)
+    os.makedirs(args.outdir, exist_ok=True)
+    with open(os.path.join(args.outdir, f"{args.tag}.json"), "w") as f:
+        json.dump(out, f)
+    print(f"saved -> {args.outdir}/{args.tag}.json", flush=True)
+
+
+if __name__ == "__main__":
+    main()