1import argparse
  2import json
  3import os
  4import time
  5import numpy as np
  6import torch
  7from torch import nn
  8from torchvision import datasets, transforms
  9
 10
 11def operator_check(d=24, gamma=0.35, lam=0.4, trials=300000, seed=17):
 12    rng = np.random.default_rng(seed)
 13    eig = np.linspace(0.03, 1.0, d) / d
 14    idx = rng.integers(0, d, trials)
 15    scale2 = d * eig[idx]
 16    p = np.full((trials, d), 1.0 - gamma * lam)
 17    p[np.arange(trials), idx] -= gamma * scale2
 18    ep2 = np.mean(p * p, axis=0)
 19    A = 1.0 - gamma * (eig + lam)
 20    rhs1 = (1.0 - gamma * lam) * A
 21    rhs2 = (1.0 - gamma * lam) ** 2
 22    gap1 = float(np.max(ep2 - rhs1))
 23    gap2 = float(np.max(rhs1 - rhs2))
 24    return {"max_E_P2_minus_(1-gamma-lambda)A": gap1,
 25            "max_(1-gamma-lambda)A_minus_scalar_bound": gap2,
 26            "gamma_lambda": gamma * lam, "A_min": float(A.min()),
 27            "A_max": float(A.max()),
 28            "passed": bool(gap1 <= 0.01 and gap2 <= 1e-12 and 0 <= gamma*lam <= 1)}
 29
 30
 31def load_mnist(root):
 32    tf = transforms.Compose([transforms.ToTensor(),
 33        transforms.Normalize((0.1307,), (0.3081,)), transforms.Lambda(lambda x: x.view(-1))])
 34    tr = datasets.MNIST(root, train=True, download=True, transform=tf)
 35    te = datasets.MNIST(root, train=False, download=True, transform=tf)
 36    # Materialize once: direct tensor indexing is substantially faster than repeated PIL transforms.
 37    Xtr = torch.stack([tr[i][0] for i in range(len(tr))])
 38    ytr = torch.tensor(tr.targets, dtype=torch.long)
 39    Xte = torch.stack([te[i][0] for i in range(len(te))])
 40    yte = torch.tensor(te.targets, dtype=torch.long)
 41    return Xtr, ytr, Xte, yte
 42
 43
 44def accuracy(model, x, y, device):
 45    model.eval(); correct = 0
 46    with torch.inference_mode():
 47        for i in range(0, len(x), 512):
 48            correct += int((model(x[i:i+512].to(device)).argmax(1).cpu() == y[i:i+512]).sum())
 49    model.train()
 50    return correct / len(x)
 51
 52
 53def train_one(data, seed, mode, device, T=3000, width=128, batch=128, wd=.05, p=.5, q=.2):
 54    Xtr, ytr, Xte, yte = data
 55    torch.manual_seed(seed); np.random.seed(seed)
 56    model = nn.Sequential(nn.Linear(784, width), nn.ReLU(), nn.Linear(width, width),
 57                          nn.ReLU(), nn.Linear(width, 10)).to(device)
 58    opt = torch.optim.AdamW(model.parameters(), lr=1e-3, betas=(.9, .999), weight_decay=wd)
 59    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=T)
 60    tail_start, tail_end = int((1-q)*T), T
 61    avg = [torch.zeros_like(x, dtype=torch.float32, device='cpu') for x in model.parameters()]
 62    avg_n = 0; best = -1.; best_step = 0; grad_mid = 0.; grad_n = 0
 63    rng = torch.Generator().manual_seed(seed + 10000)
 64    t0 = time.perf_counter()
 65    for step in range(T):
 66        ix = torch.randint(len(Xtr), (batch,), generator=rng)
 67        xb, yb = Xtr[ix].to(device), ytr[ix].to(device)
 68        opt.param_groups[0]['weight_decay'] = wd if (mode in ('baseline','wd_tail') or step < int(p*T)) else 0.
 69        opt.zero_grad(set_to_none=True); loss = nn.functional.cross_entropy(model(xb), yb); loss.backward()
 70        gn = torch.sqrt(sum((x.grad.detach()**2).sum() for x in model.parameters() if x.grad is not None)).item()
 71        if 2000 <= step < 3000: grad_mid += gn; grad_n += 1
 72        opt.step(); sched.step()
 73        if mode.endswith('tail') and tail_start <= step < tail_end:
 74            with torch.no_grad():
 75                for a, x in zip(avg, model.parameters()): a.add_(x.detach().cpu())
 76            avg_n += 1
 77        # Test every 300 steps plus final: adequate for a small deterministic super-check.
 78        if (step + 1) % 300 == 0 or step == T-1:
 79            acc = accuracy(model, Xte, yte, device)
 80            if acc > best: best, best_step = acc, step + 1
 81    last_acc = acc; tail_acc = None
 82    if mode.endswith('tail') and avg_n:
 83        saved = [x.detach().clone() for x in model.parameters()]
 84        with torch.no_grad():
 85            for x, a in zip(model.parameters(), avg): x.copy_((a / avg_n).to(device))
 86        tail_acc = accuracy(model, Xte, yte, device)
 87        with torch.no_grad():
 88            for x, s in zip(model.parameters(), saved): x.copy_(s)
 89    norm = torch.sqrt(sum((x.detach()**2).sum() for x in model.parameters())).item()
 90    return {'test_accuracy': float(tail_acc if tail_acc is not None else last_acc),
 91            'last_nonaveraged_accuracy': float(last_acc), 'best_test_accuracy': float(best),
 92            'steps_to_best': int(best_step), 'final_weight_norm': float(norm),
 93            'mid_gradient_norm': float(grad_mid/max(1,grad_n)), 'runtime_sec': time.perf_counter()-t0}
 94
 95
 96def main():
 97    ap = argparse.ArgumentParser(); ap.add_argument('--runs', type=int, default=5)
 98    ap.add_argument('--T', type=int, default=3000); ap.add_argument('--width', type=int, default=128)
 99    ap.add_argument('--out', default='super_results.json'); args = ap.parse_args()
100    try: device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
101    except Exception: device = torch.device('cpu')
102    root = os.path.join(os.path.dirname(__file__), 'data'); data = load_mnist(root)
103    check = operator_check(); modes = ['baseline','initwd','initwd_tail','wd_tail']; raw={m:[] for m in modes}; start=time.perf_counter()
104    for seed in range(args.runs):
105        for mode in modes:
106            try: row=train_one(data, seed, mode, device, T=args.T, width=args.width)
107            except RuntimeError as e:
108                if device.type=='cuda':
109                    torch.cuda.empty_cache(); device=torch.device('cpu'); row=train_one(data, seed, mode, device, T=args.T, width=args.width)
110                else: raise
111            raw[mode].append(row); print(mode, seed, row, flush=True)
112    summary={}
113    for mode, rows in raw.items():
114        summary[mode]={}
115        for key in ['test_accuracy','best_test_accuracy','steps_to_best','final_weight_norm','mid_gradient_norm','runtime_sec']:
116            v=np.array([r[key] for r in rows],float); summary[mode][key+'_mean']=float(v.mean()); summary[mode][key+'_std']=float(v.std(ddof=1))
117    out={'config':{'dataset':'MNIST train=60000 test=10000','T':args.T,'width':args.width,'batch':128,'lr':1e-3,'betas':[.9,.999],'weight_decay':.05,'p':.5,'q':.2,'seeds':args.runs,'device':str(device),'total_runtime_sec':time.perf_counter()-start},'operator_check':check,'summary':summary,'raw':raw}
118    with open(args.out,'w') as f: json.dump(out,f,indent=2)
119    print(json.dumps({'config':out['config'],'operator_check':check,'summary':summary},indent=2))
120
121if __name__ == '__main__': main()