1import sys, os, json, math, random, copy, time
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  7from bench import (get_dataset, make_model, train_model, evaluate,
  8                   sweep_baseline, make_report)
  9
 10SEEDS = tuple(range(8))
 11# Shared union of learning rates: every idea lr is also a baseline lr.
 12GRID = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 1e-2}]
 13EPOCHS = 8
 14BATCH = 128
 15NTRAIN, NTEST = 400, 200
 16
 17
 18def seed_all(seed):
 19    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 20    if torch.cuda.is_available():
 21        torch.cuda.manual_seed_all(seed)
 22    torch.set_num_threads(4)
 23
 24
 25def dataset(seed):
 26    return get_dataset("tabular", int(seed), n_train=NTRAIN, n_test=NTEST)
 27
 28
 29def baseline_fn(cfg):
 30    def run(seed):
 31        seed_all(seed)
 32        d = dataset(seed)
 33        seed_all(seed + 10000)
 34        net = make_model("mlp_tiny", d["input_shape"], d["out_dim"])
 35        _, metric, _ = train_model(net, d, epochs=EPOCHS, lr=cfg["lr"],
 36                                   batch=BATCH, weight_decay=0.0, log=lambda *_: None)
 37        return float(metric)
 38    return run
 39
 40
 41def loss_fn(net, x, y):
 42    return ((net(x) - y) ** 2).mean()
 43
 44
 45def apply_step(net, grads, eta):
 46    with torch.no_grad():
 47        for p, g in zip(net.parameters(), grads):
 48            p.add_(g, alpha=-float(eta))
 49
 50
 51def controller_train(net, d, initial_lr, alpha=1.0, rho=0.9, c=0.1,
 52                     beta=0.5, eta_min=1e-5, eta_max=1.0, max_bt=5):
 53    # This is deliberately a training loop intervention: probe and candidate
 54    # losses are evaluated on the same minibatch as F0.
 55    device = next(net.parameters()).device
 56    xtr, ytr = d["xtr"].to(device), d["ytr"].to(device)
 57    eta_probe = float(initial_lr)
 58    accepted = 0; total = 0; observed_lh = []; chosen = []
 59    for ep in range(EPOCHS):
 60        net.train()
 61        perm = torch.randperm(len(xtr), device=device)
 62        for i in range(0, len(xtr), BATCH):
 63            idx = perm[i:i+BATCH]; x, y = xtr[idx], ytr[idx]
 64            net.zero_grad(set_to_none=True)
 65            f0 = loss_fn(net, x, y)
 66            grads = torch.autograd.grad(f0, tuple(net.parameters()))
 67            g2 = sum(float((g.detach() ** 2).sum()) for g in grads)
 68            G = math.sqrt(max(g2, 0.0))
 69            total += 1
 70            if G < 1e-12:
 71                continue
 72            # Probe at maintained rate, then restore exactly.
 73            apply_step(net, grads, eta_probe)
 74            with torch.no_grad():
 75                fp = float(loss_fn(net, x, y))
 76            apply_step(net, grads, -eta_probe)
 77            snorm = eta_probe * G
 78            rem = fp - float(f0.detach()) + eta_probe * g2
 79            lhat = (1.0 + alpha) * max(rem, 0.0) / (snorm ** (1.0 + alpha) + 1e-12)
 80            observed_lh.append(lhat)
 81            eta_star = (G ** (1.0-alpha) / (lhat + 1e-12)) ** (1.0/alpha)
 82            eta = min(eta_max, max(eta_min, rho * eta_star))
 83            ok = False
 84            fchosen = None
 85            for _ in range(max_bt + 1):
 86                apply_step(net, grads, eta)
 87                with torch.no_grad():
 88                    fc = float(loss_fn(net, x, y))
 89                if fc <= float(f0.detach()) - c * eta * g2 + 1e-10:
 90                    ok = True; fchosen = fc; break
 91                apply_step(net, grads, -eta)
 92                eta *= beta
 93            if not ok:
 94                # candidate is currently applied on final failed attempt
 95                apply_step(net, grads, -eta)
 96                eta_probe = max(eta_min, eta_probe * beta)
 97            else:
 98                accepted += 1
 99                chosen.append(eta)
100                # next probe gently grows from accepted scale
101                eta_probe = min(eta_max, eta / 0.9)
102    net.eval()
103    with torch.no_grad():
104        xte, yte = d["xte"].to(device), d["yte"].to(device)
105        metric = float(loss_fn(net, xte, yte))
106    stats = {"accepted_fraction": accepted / max(total, 1),
107             "mean_lhat": float(np.mean(observed_lh)) if observed_lh else 0.0,
108             "mean_accepted_eta": float(np.mean(chosen)) if chosen else 0.0,
109             "n_updates": total}
110    return metric, stats
111
112
113def idea_fn(cfg, collect=False):
114    def run(seed):
115        seed_all(seed)
116        d = dataset(seed)
117        seed_all(seed + 10000)
118        net = make_model("mlp_tiny", d["input_shape"], d["out_dim"])
119        # Use official fallback placement convention by trying CUDA, then CPU.
120        device = "cuda" if torch.cuda.is_available() else "cpu"
121        try:
122            net = net.to(device)
123            metric, stats = controller_train(net, d, cfg["lr"])
124        except RuntimeError:
125            net = make_model("mlp_tiny", d["input_shape"], d["out_dim"]).to("cpu")
126            metric, stats = controller_train(net, d, cfg["lr"])
127        if collect:
128            RUN_STATS.setdefault(str(cfg["lr"]), {})[str(seed)] = stats
129        return float(metric)
130    return run
131
132RUN_STATS = {}
133if __name__ == "__main__":
134    t0 = time.time()
135    base = sweep_baseline(baseline_fn, GRID)
136    # Idea sweep uses exactly the same three learning rates; select on sweep seeds.
137    idea_sweep = []
138    for cfg in GRID:
139        r = evaluate(idea_fn(cfg), seeds=(0,1,2,3))
140        idea_sweep.append({"cfg": cfg, "mean": r["mean"]})
141    best_cfg = min(GRID, key=lambda c: next(x["mean"] for x in idea_sweep if x["cfg"] == c))
142    idea_full = evaluate(idea_fn(best_cfg, collect=True), seeds=SEEDS)
143    # Re-test the NN-scale quantitative prediction: observed Lhat should predict
144    # the accepted scale eta approximately as rho/Lhat for alpha=1.
145    vals = []
146    for st in RUN_STATS.get(str(best_cfg["lr"]), {}).values():
147        if st["mean_lhat"] > 0 and st["mean_accepted_eta"] > 0:
148            pred = 0.9 / st["mean_lhat"]
149            obs = st["mean_accepted_eta"]
150            vals.append({"predicted_eta": pred, "observed_eta": obs,
151                         "ratio_observed_to_predicted": obs/pred})
152    ratios = [v["ratio_observed_to_predicted"] for v in vals]
153    signature = {"alpha": 1.0, "rho": 0.9, "prediction": "accepted_eta ~= rho / observed_directional_Lhat",
154                 "per_seed": vals,
155                 "mean_ratio": float(np.mean(ratios)) if ratios else None,
156                 "confirmed": bool(ratios and 0.5 <= float(np.mean(ratios)) <= 1.5)}
157    base["idea_parity_sweep"] = idea_sweep
158    report = make_report("tabular", "mlp_tiny", base, idea_full, signature)
159    report["idea"]["selected_cfg"] = best_cfg
160    report["runtime_seconds"] = time.time() - t0
161    report["protocol_notes"] = {"epochs": EPOCHS, "batch": BATCH,
162        "dataset_sizes": [NTRAIN, NTEST], "structural_match": "optimizer -> tabular"}
163    with open("bench_report.json", "w") as f: json.dump(report, f, indent=2)
164    print(json.dumps(report, indent=2))