diff options
| -rw-r--r-- | experiments/run.py | 46 |
1 files changed, 36 insertions, 10 deletions
diff --git a/experiments/run.py b/experiments/run.py index bc66bae..0ace433 100644 --- a/experiments/run.py +++ b/experiments/run.py @@ -81,6 +81,21 @@ def hardware_report(device): return report +def fixed_training_probe(train_loader, probe_bs, device): + """Return a deterministic diagnostic batch without advancing data order. + + All project loaders expose their in-memory tensors. Slicing those tensors + avoids consuming the loader's shuffle generator and, importantly, keeps a + held-out validation/test split out of gradient-alignment diagnostics. + """ + if not hasattr(train_loader, "x") or not hasattr(train_loader, "y"): + raise TypeError("diagnostic probes require an in-memory training loader") + n = min(probe_bs, len(train_loader.x)) + if n <= 0: + raise ValueError("diagnostic probe requires at least one training example") + return train_loader.x[:n].to(device), train_loader.y[:n].to(device) + + def calibration_work_per_event(net, cfg): """Hardware-independent causal-calibration work for one minibatch. @@ -224,16 +239,19 @@ def train(args): # resident training state plus transient perturbation-batch expansion. reset_peak_memory(device) - # A fixed probe batch is loaded only when diagnostics are requested. In a - # frozen test-only run with diagnostics disabled, the test set is therefore - # not touched until the single final evaluation. + # Diagnostics use a fixed training prefix without advancing the shuffled + # training iterator. Held-out data are reserved for metric evaluation. px = py = poh = None if args.diagnostics != "none": - px, py = next(iter(eval_loader)) - px, py = px[:args.probe_bs].to(device), py[:args.probe_bs].to(device) + px, py = fixed_training_probe(train_loader, args.probe_bs, device) poh = onehot(py, n_out, device=device) log = {"args": vars(args), "split": split, "provenance": code_provenance(), + "diagnostic_protocol": { + "probe_source": "training_prefix" if px is not None else None, + "probe_examples": len(px) if px is not None else 0, + "schedule": args.diagnostics_schedule, + }, "steps": [], "final": {}} step = 0 prev_error = None @@ -298,12 +316,14 @@ def train(args): if step % args.log_every == 0: rec = {"step": step, "epoch": epoch, "train_loss": float(loss)} - if args.diagnostics != "none": + inline_diagnostics = (args.diagnostics != "none" + and args.diagnostics_schedule == "inline") + if inline_diagnostics: device_sync(device) diagnostics_t0 = time.time() - if args.mode == "fa" and args.diagnostics != "none": + if args.mode == "fa" and inline_diagnostics: rec.update(probes.fa_alignment_report(net, px, py, poh)) - elif args.mode != "bp" and args.diagnostics != "none": + elif args.mode != "bp" and inline_diagnostics: al = probes.alignment_report(net, px, py, poh, cfg) rec["cos_r_negg"] = al["cos_r_negg"] rec["cos_innovation_negg"] = al["cos_innovation_negg"] @@ -316,7 +336,7 @@ def train(args): if (args.mode == "sdil" and args.diagnostics == "full" and step % (args.log_every * 5) == 0): rec["ldr"] = probes.loss_decrease_ratio(net, px, py, poh, cfg, step) - if args.diagnostics != "none": + if inline_diagnostics: device_sync(device) elapsed = time.time() - diagnostics_t0 diagnostics_wall_s += elapsed @@ -350,7 +370,8 @@ def train(args): record.update({"test_acc": acc, "test_loss": tloss}) else: record.update({"val_acc": acc, "val_loss": tloss}) - if args.diagnostics != "none": + if (args.diagnostics != "none" + and args.diagnostics_schedule == "inline"): device_sync(device) diagnostics_t0 = time.time() if args.mode == "fa": @@ -391,6 +412,8 @@ def train(args): log["final"]["cos_innovation_negg"] = al["cos_innovation_negg"] log["final"]["cos_apical_negg"] = al["cos_apical_negg"] log["final"]["cos_Ac_negg"] = al["cos_Ac_negg"] + log["final"]["r_norm"] = al["r_norm"] + log["final"]["g_norm"] = al["g_norm"] log["final"]["traffic_norm"] = al["traffic_norm"] log["final"]["traffic_residual_norm"] = al["traffic_residual_norm"] log["final"]["traffic_r2"] = al["traffic_r2"] @@ -489,6 +512,9 @@ def get_args(): p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") p.add_argument("--log_every", type=int, default=50) p.add_argument("--diagnostics", default="full", choices=["none", "alignment", "full"]) + p.add_argument("--diagnostics_schedule", default="inline", + choices=["inline", "final"], + help="run probes throughout training or only after final evaluation") p.add_argument("--eval_every", type=int, default=1, help="0 evaluates only once at the end") p.add_argument("--max_steps", type=int, default=0) # 0 = no cap (smoke only) |