diff options
Diffstat (limited to 'experiments/run.py')
| -rw-r--r-- | experiments/run.py | 191 |
1 files changed, 163 insertions, 28 deletions
diff --git a/experiments/run.py b/experiments/run.py index 8c59a50..aa509e4 100644 --- a/experiments/run.py +++ b/experiments/run.py @@ -48,6 +48,40 @@ def code_provenance(): return {"git_commit": None, "git_dirty": None} +def device_sync(device): + if str(device).startswith("cuda") and torch.cuda.is_available(): + torch.cuda.synchronize() + + +def calibration_work_per_event(net, cfg): + """Hardware-independent causal-calibration work for one minibatch. + + Forward equivalents use affine multiply-add work as the denominator. The + returned loss evaluations count calls producing per-example scalar losses; + the caller multiplies by the actual minibatch size for scalar observations. + """ + directions = cfg.pert_ndirs + if cfg.pert_mode == "simultaneous": + return { + "batch_loss_evaluations": 2 * directions, + # The implementation performs one duplicate clean forward plus one + # batched +/- forward for every direction. + "forward_equivalent_batches": 1.0 + 2.0 * directions, + "perturbation_batch_expansion": 2 * directions, + } + + full_work = sum(weight.numel() for weight in net.W) + tail_work = 0 + for hidden_layer in range(net.L - 1): + tail_work += sum(net.W[j].numel() for j in range(hidden_layer + 1, net.L)) + return { + # One clean baseline loss plus +/- tail loss for every hidden layer. + "batch_loss_evaluations": 1 + 2 * directions * (net.L - 1), + "forward_equivalent_batches": 1.0 + 2.0 * directions * tail_work / full_work, + "perturbation_batch_expansion": 1, + } + + def build(args, device): sizes = [args.n_in] + [args.width] * args.depth + [args.n_out] if args.mode == "bp": @@ -159,20 +193,38 @@ def train(args): args.n_in, args.n_out = n_in, n_out net, cfg = build(args, device) - # a fixed probe batch for stable alignment tracking - px, py = next(iter(eval_loader)) - px, py = px[:args.probe_bs].to(device), py[:args.probe_bs].to(device) - poh = onehot(py, n_out, device=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. + 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) + poh = onehot(py, n_out, device=device) log = {"args": vars(args), "split": split, "provenance": code_provenance(), "steps": [], "final": {}} step = 0 prev_error = None + device_sync(device) t0 = time.time() + train_wall_s = 0.0 + eval_wall_s = 0.0 + diagnostics_wall_s = 0.0 + inline_diagnostics_wall_s = 0.0 + warmup_examples = 0 + ordinary_forward_examples = 0 + perturbation_events = 0 + calibration_batch_loss_evaluations = 0 + calibration_example_loss_evaluations = 0 + calibration_forward_equivalent_examples = 0.0 + perturbation_batch_expansion = 0 # predictor warmup on neutral-period (c=0) drive, so P cancels the apical # nuisance before task plasticity relies on the residual (no-op when rho=0). if args.mode == "sdil" and args.learn_P and args.p_warmup_steps > 0 and args.nuis_rho > 0: + device_sync(device) + warmup_t0 = time.time() it = iter(train_loader) for _ in range(args.p_warmup_steps): try: @@ -180,10 +232,19 @@ def train(args): except StopIteration: it = iter(train_loader) wx, _ = next(it) - neutral_p_update(net, wx.to(device), args.p_warmup_eta) + wx = wx.to(device) + warmup_examples += wx.shape[0] + neutral_p_update(net, wx, args.p_warmup_eta) + device_sync(device) + warmup_wall_s = time.time() - warmup_t0 + else: + warmup_wall_s = 0.0 for epoch in range(args.epochs): + device_sync(device) + train_t0 = time.time() for x, y in train_loader: x, y = x.to(device), y.to(device) + ordinary_forward_examples += x.shape[0] yoh = onehot(y, n_out, device=device) if args.mode == "bp": loss = net.bp_step(x, y, cfg.eta, momentum=cfg.momentum) @@ -192,12 +253,26 @@ def train(args): else: loss, aux = sdil_step(net, x, y, yoh, cfg, step, prev_error=prev_error) prev_error = aux["error"] + if args.mode == "sdil" and aux["did_pert"]: + event = calibration_work_per_event(net, cfg) + perturbation_events += 1 + calibration_batch_loss_evaluations += event["batch_loss_evaluations"] + calibration_example_loss_evaluations += ( + event["batch_loss_evaluations"] * x.shape[0]) + calibration_forward_equivalent_examples += ( + event["forward_equivalent_batches"] * x.shape[0]) + perturbation_batch_expansion = max( + perturbation_batch_expansion, + event["perturbation_batch_expansion"]) if step % args.log_every == 0: rec = {"step": step, "epoch": epoch, "train_loss": float(loss)} - if args.mode == "fa": + if args.diagnostics != "none": + device_sync(device) + diagnostics_t0 = time.time() + if args.mode == "fa" and args.diagnostics != "none": rec.update(probes.fa_alignment_report(net, px, py, poh)) - elif args.mode != "bp": + elif args.mode != "bp" and args.diagnostics != "none": 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"] @@ -207,45 +282,79 @@ def train(args): rec["traffic_norm"] = al["traffic_norm"] rec["traffic_residual_norm"] = al["traffic_residual_norm"] rec["traffic_r2"] = al["traffic_r2"] - if args.mode == "sdil" and step % (args.log_every * 5) == 0: + 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": + device_sync(device) + elapsed = time.time() - diagnostics_t0 + diagnostics_wall_s += elapsed + inline_diagnostics_wall_s += elapsed log["steps"].append(rec) step += 1 if args.max_steps and step >= args.max_steps: break if args.max_steps and step >= args.max_steps: + device_sync(device) + train_wall_s += time.time() - train_t0 break - acc, tloss = evaluate(net, eval_loader) - metric = "val_acc" if args.eval_split == "validation" else "test_acc" - msg = f"[{args.tag}] epoch {epoch} step {step} loss {loss:.4f} {metric} {acc:.4f}" - if args.mode == "fa": - al = probes.fa_alignment_report(net, px, py, poh) - meancos = sum(al["cos_fa_negg"]) / len(al["cos_fa_negg"]) - msg += f" mean_cos(fa,-g) {meancos:+.3f} per-layer {['%.2f'%v for v in al['cos_fa_negg']]}" - elif args.mode != "bp": - al = probes.alignment_report(net, px, py, poh, cfg) - meancos = sum(al["cos_r_negg"]) / len(al["cos_r_negg"]) - msg += f" mean_cos(r,-g) {meancos:+.3f} per-layer {['%.2f'%v for v in al['cos_r_negg']]}" + device_sync(device) + train_wall_s += time.time() - train_t0 + + should_evaluate = args.eval_every > 0 and (epoch + 1) % args.eval_every == 0 + record = {"epoch_end": epoch, "step": step} + msg = f"[{args.tag}] epoch {epoch} step {step} loss {loss:.4f}" + if should_evaluate: + device_sync(device) + eval_t0 = time.time() + acc, tloss = evaluate(net, eval_loader) + device_sync(device) + eval_wall_s += time.time() - eval_t0 + metric = "val_acc" if args.eval_split == "validation" else "test_acc" + msg += f" {metric} {acc:.4f}" + record.update({"eval_split": args.eval_split, + "eval_acc": acc, "eval_loss": tloss}) + if args.eval_split == "test": + record.update({"test_acc": acc, "test_loss": tloss}) + else: + record.update({"val_acc": acc, "val_loss": tloss}) + if args.diagnostics != "none": + device_sync(device) + diagnostics_t0 = time.time() + if args.mode == "fa": + al = probes.fa_alignment_report(net, px, py, poh) + meancos = sum(al["cos_fa_negg"]) / len(al["cos_fa_negg"]) + msg += f" mean_cos(fa,-g) {meancos:+.3f} per-layer {['%.2f'%v for v in al['cos_fa_negg']]}" + elif args.mode != "bp": + al = probes.alignment_report(net, px, py, poh, cfg) + meancos = sum(al["cos_r_negg"]) / len(al["cos_r_negg"]) + msg += f" mean_cos(r,-g) {meancos:+.3f} per-layer {['%.2f'%v for v in al['cos_r_negg']]}" + device_sync(device) + diagnostics_wall_s += time.time() - diagnostics_t0 print(msg, flush=True) - record = {"epoch_end": epoch, "step": step, "eval_split": args.eval_split, - "eval_acc": acc, "eval_loss": tloss} - if args.eval_split == "test": - record.update({"test_acc": acc, "test_loss": tloss}) - else: - record.update({"val_acc": acc, "val_loss": tloss}) log["steps"].append(record) + device_sync(device) + eval_t0 = time.time() acc, tloss = evaluate(net, eval_loader) + device_sync(device) + eval_wall_s += time.time() - eval_t0 log["final"] = {"eval_split": args.eval_split, "eval_acc": acc, "eval_loss": tloss, - "wall_s": time.time() - t0} + "wall_s": None} if args.eval_split == "test": log["final"].update({"test_acc": acc, "test_loss": tloss}) else: log["final"].update({"val_acc": acc, "val_loss": tloss}) - if args.mode == "fa": + if args.mode == "fa" and args.diagnostics != "none": + device_sync(device) + diagnostics_t0 = time.time() log["final"].update(probes.fa_alignment_report(net, px, py, poh)) - elif args.mode != "bp": + device_sync(device) + diagnostics_wall_s += time.time() - diagnostics_t0 + elif args.mode != "bp" and args.diagnostics != "none": + device_sync(device) + diagnostics_t0 = time.time() al = probes.alignment_report(net, px, py, poh, cfg) log["final"]["cos_r_negg"] = al["cos_r_negg"] log["final"]["cos_innovation_negg"] = al["cos_innovation_negg"] @@ -254,6 +363,29 @@ def train(args): log["final"]["traffic_norm"] = al["traffic_norm"] log["final"]["traffic_residual_norm"] = al["traffic_residual_norm"] log["final"]["traffic_r2"] = al["traffic_r2"] + device_sync(device) + diagnostics_wall_s += time.time() - diagnostics_t0 + device_sync(device) + log["final"]["wall_s"] = time.time() - t0 + log["timing"] = { + "warmup_wall_s": warmup_wall_s, + "training_loop_wall_s": train_wall_s, + "diagnostics_wall_s": diagnostics_wall_s, + "inline_diagnostics_wall_s": inline_diagnostics_wall_s, + "optimizer_wall_s_excluding_inline_diagnostics": max( + 0.0, train_wall_s - inline_diagnostics_wall_s), + "evaluation_wall_s": eval_wall_s, + } + log["cost"] = { + "train_steps": step, + "ordinary_training_forward_examples": ordinary_forward_examples, + "predictor_warmup_forward_examples": warmup_examples, + "perturbation_events": perturbation_events, + "calibration_batch_loss_evaluations": calibration_batch_loss_evaluations, + "calibration_example_loss_evaluations": calibration_example_loss_evaluations, + "calibration_forward_equivalent_examples": calibration_forward_equivalent_examples, + "max_perturbation_batch_expansion": perturbation_batch_expansion, + } os.makedirs(args.outdir, exist_ok=True) outpath = os.path.join(args.outdir, f"{args.tag}.json") with open(outpath, "w") as f: @@ -319,6 +451,9 @@ def get_args(): p.add_argument("--seed", type=int, default=0) 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("--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) p.add_argument("--probe_bs", type=int, default=512) p.add_argument("--outdir", default="results") |