experiments: isolate final diagnostics from held-out data

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)