diff options
Diffstat (limited to 'experiments/baseline_run.py')
| -rw-r--r-- | experiments/baseline_run.py | 136 |
1 files changed, 114 insertions, 22 deletions
diff --git a/experiments/baseline_run.py b/experiments/baseline_run.py index a96e24d..24ad6ae 100644 --- a/experiments/baseline_run.py +++ b/experiments/baseline_run.py @@ -19,6 +19,7 @@ import torch.nn.functional as F sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from sdil.data import get_dataset, onehot from sdil.local_baselines import FANet, PEPITANet, FFNet, EPNet +from sdil import probes METHOD_SOURCES = { @@ -39,7 +40,7 @@ METHOD_SOURCES = { "ep": { "paper": "https://www.frontiersin.org/journals/computational-neuroscience/articles/10.3389/fncom.2017.00024/full", "code": "https://github.com/bscellier/Towards-a-Biologically-Plausible-Backprop", - "protocol": "two-phase energy-gradient dynamics; random beta sign", + "protocol": "two-phase energy-gradient dynamics; persistent free particles; random beta sign", }, } @@ -108,10 +109,54 @@ def ep_learning_rates(depth, base_eta=None): return [base_eta / (4 ** l) for l in range(depth + 1)] if depth in canonical: return canonical[depth] - # The original implementation already required rapidly shrinking rates as - # depth grew. Continue its depth-3 geometric schedule for an explicit, - # logged extrapolation rather than pretending there is a canonical one. - return [0.128 / (4 ** l) for l in range(depth + 1)] + raise ValueError("non-canonical EP depth requires an explicit --eta") + + +def ep_dynamics(depth, beta=None, free_steps=None, nudge_steps=None): + """Original-model dynamics for the published 1/2/3-hidden-layer runs.""" + canonical = { + 1: {"beta": 0.5, "free_steps": 20, "nudge_steps": 4}, + 2: {"beta": 1.0, "free_steps": 150, "nudge_steps": 6}, + 3: {"beta": 1.0, "free_steps": 500, "nudge_steps": 8}, + } + if depth not in canonical and (beta is None or free_steps is None or nudge_steps is None): + raise ValueError( + "non-canonical EP depth requires --ep_beta, --ep_free_steps, and --ep_nudge_steps") + defaults = canonical.get(depth, {}) + return { + "beta": defaults["beta"] if beta is None else beta, + "free_steps": defaults["free_steps"] if free_steps is None else free_steps, + "nudge_steps": defaults["nudge_steps"] if nudge_steps is None else nudge_steps, + "canonical_depth": depth in canonical, + } + + +def pepita_protocol(args): + """Resolve only settings that are actually supported by a cited protocol. + + The original PEPITA paper specifies one hidden layer. Deep PEPITA is very + sensitive to architecture-specific learning rates, so silently reusing the + shallow 0.1 default is worse than requiring an explicit value. + """ + eta = args.eta + if eta is None: + if args.depth != 1: + raise ValueError("deep PEPITA requires an explicit --eta") + if args.dataset == "mnist": + eta = 0.1 + elif args.dataset == "cifar10": + eta = 0.01 + else: + raise ValueError("PEPITA on this dataset requires an explicit --eta") + if args.pepita_decay_epochs is not None: + decay_epochs = [int(v) for v in args.pepita_decay_epochs.split(",") if v.strip()] + elif args.depth == 1 and args.dataset == "mnist": + decay_epochs = [60] + elif args.depth == 1 and args.dataset == "cifar10": + decay_epochs = [60, 90] + else: + decay_epochs = [] + return {"eta": eta, "decay_epochs": decay_epochs} def build(args, n_in, device): @@ -128,30 +173,53 @@ def build(args, n_in, device): return FFNet([n_in] + [args.width] * args.depth, device=device, seed=args.seed, threshold=args.ff_threshold) if args.method == "ep": - return EPNet(sizes, device=device, seed=args.seed, beta=args.ep_beta, - dt=args.ep_dt, T_free=args.ep_free_steps, - T_nudge=args.ep_nudge_steps, random_beta_sign=True) + dyn = ep_dynamics(args.depth, args.ep_beta, args.ep_free_steps, + args.ep_nudge_steps) + return EPNet(sizes, device=device, seed=args.seed, beta=dyn["beta"], + dt=args.ep_dt, T_free=dyn["free_steps"], + T_nudge=dyn["nudge_steps"], random_beta_sign=True) raise ValueError(args.method) def train(args): torch.manual_seed(args.seed) device = args.device + pepita_cfg = pepita_protocol(args) if args.method == "pepita" else None + ep_dyn = (ep_dynamics(args.depth, args.ep_beta, args.ep_free_steps, + args.ep_nudge_steps) if args.method == "ep" else None) + use_persistent_ep = args.method == "ep" and bool(args.ep_persistent) train_loader, test_loader, n_in, n_out = get_dataset( - args.dataset, args.batch_size, device=device) + args.dataset, args.batch_size, device=device, + shuffle_train=not use_persistent_ep, + train_limit=args.train_examples or None) net = build(args, n_in, device) canonical = bool(args.canonical_preprocess) + resolved_protocol = {} + if pepita_cfg is not None: + resolved_protocol.update(pepita_cfg) + if ep_dyn is not None: + resolved_protocol.update(ep_dyn) + resolved_protocol["learning_rates"] = ep_learning_rates(args.depth, args.eta) + resolved_protocol["persistent_particles"] = use_persistent_ep + resolved_protocol["train_examples"] = train_loader.n log = { "args": vars(args), + "resolved_protocol": resolved_protocol, "method_source": METHOD_SOURCES[args.method], "provenance": code_provenance(), "steps": [], "final": {}, } + if args.method == "fa": + px, py = next(iter(test_loader)) + px, py = px[:args.probe_bs], py[:args.probe_bs] + pyoh = onehot(py, n_out, device=device) t0 = time.time() if args.method == "ff": eta = 0.03 if args.eta is None else args.eta + resolved_protocol["eta"] = eta + resolved_protocol["epochs_per_layer"] = args.epochs for layer in range(args.depth): for epoch in range(args.epochs): batches = 0 @@ -168,15 +236,21 @@ def train(args): "local_loss": loss}) print(f"[{args.tag}] layer {layer} test_acc {acc:.4f}", flush=True) else: - eta = args.eta - if eta is None: - eta = 0.05 if args.method == "fa" else (0.1 if args.method == "pepita" else None) + if args.method == "fa": + eta = 0.05 if args.eta is None else args.eta + resolved_protocol["eta"] = eta + resolved_protocol["feedback_scale"] = args.feedback_scale + elif args.method == "pepita": + eta = pepita_cfg["eta"] + else: + eta = args.eta ep_etas = ep_learning_rates(args.depth, eta) if args.method == "ep" else None + ep_states = [None] * len(train_loader) if use_persistent_ep else None for epoch in range(args.epochs): - if args.method == "pepita" and epoch in (60, 90): + if args.method == "pepita" and epoch in pepita_cfg["decay_epochs"]: eta *= 0.1 batches = 0 - for x, y in train_loader: + for batch_index, (x, y) in enumerate(train_loader): x = canonical_input(x, args.dataset, args.method, canonical) yoh = onehot(y, n_out, device=device) if args.method == "fa": @@ -184,21 +258,33 @@ def train(args): elif args.method == "pepita": loss = net.pepita_step(x, y, yoh, eta, args.momentum) else: - loss = net.train_step(x, y, yoh, ep_etas) + if use_persistent_ep: + loss, ep_states[batch_index] = net.train_step( + x, y, yoh, ep_etas, free_state=ep_states[batch_index], + return_free_state=True) + else: + loss = net.train_step(x, y, yoh, ep_etas) batches += 1 if args.max_batches and batches >= args.max_batches: break acc, test_loss = evaluate_method(net, test_loader, args.method, args.dataset, canonical) - log["steps"].append({"epoch_end": epoch, "test_acc": acc, - "test_loss": test_loss, "train_loss": loss}) - print(f"[{args.tag}] epoch {epoch} loss {loss:.4f} test_acc {acc:.4f}", - flush=True) + rec = {"epoch_end": epoch, "test_acc": acc, + "test_loss": test_loss, "train_loss": loss} + msg = f"[{args.tag}] epoch {epoch} loss {loss:.4f} test_acc {acc:.4f}" + if args.method == "fa": + al = probes.fa_alignment_report(net, px, py, pyoh) + rec["cos_fa_negg"] = al["cos_fa_negg"] + msg += f" mean_cos(fa,-g) {sum(al['cos_fa_negg']) / len(al['cos_fa_negg']):+.3f}" + log["steps"].append(rec) + print(msg, flush=True) acc, test_loss = evaluate_method(net, test_loader, args.method, args.dataset, canonical) log["final"] = {"test_acc": acc, "test_loss": test_loss, "wall_s": time.time() - t0} + if args.method == "fa": + log["final"].update(probes.fa_alignment_report(net, px, py, pyoh)) os.makedirs(args.outdir, exist_ok=True) path = os.path.join(args.outdir, f"{args.tag}.json") with open(path, "w") as f: @@ -216,6 +302,8 @@ def get_args(): p.add_argument("--epochs", type=int, default=10, help="FF: epochs per greedy layer; other methods: global epochs") p.add_argument("--batch_size", type=int, default=64) + p.add_argument("--train_examples", type=int, default=0, + help="0 uses the full training split") p.add_argument("--max_batches", type=int, default=0) p.add_argument("--eta", type=float, default=None) p.add_argument("--momentum", type=float, default=0.9) @@ -223,13 +311,17 @@ def get_args(): p.add_argument("--residual", type=int, default=0) p.add_argument("--feedback_scale", type=float, default=1.0) p.add_argument("--pepita_f_scale", type=float, default=0.05) + p.add_argument("--pepita_decay_epochs", default=None, + help="comma-separated; deep PEPITA defaults to no decay") p.add_argument("--keep_prob", type=float, default=0.9) p.add_argument("--ff_threshold", type=float, default=2.0) - p.add_argument("--ep_beta", type=float, default=0.5) + p.add_argument("--ep_beta", type=float, default=None) p.add_argument("--ep_dt", type=float, default=0.5) - p.add_argument("--ep_free_steps", type=int, default=20) - p.add_argument("--ep_nudge_steps", type=int, default=4) + p.add_argument("--ep_free_steps", type=int, default=None) + p.add_argument("--ep_nudge_steps", type=int, default=None) + p.add_argument("--ep_persistent", type=int, default=1) p.add_argument("--canonical_preprocess", type=int, default=1) + p.add_argument("--probe_bs", type=int, default=512) 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") |