1import sys, json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import get_dataset, make_model, sweep_baseline, make_report
  9
 10SEEDS = tuple(range(8))
 11SWEEP_SEEDS = tuple(range(4))
 12EPOCHS = 8
 13NTRAIN, NTEST = 1000, 400
 14BASE_BATCHES = (32, 64)
 15LRS = (0.0015, 0.003, 0.006)
 16
 17
 18def seed_all(seed):
 19    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 20    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 21
 22
 23def coefficients(etas):
 24    etas = np.asarray(etas, dtype=float); t = len(etas)
 25    c = np.empty(t)
 26    for i in range(t-1): c[i] = etas[i]**2 / (2*np.sum(etas[i+1:]))
 27    c[-1] = etas[-1]/2
 28    return c
 29
 30
 31def integer_allocation(weights, total, lo=8, hi=128):
 32    n = len(weights)
 33    total = int(total)
 34    lo = min(lo, total//n); hi = max(hi, lo)
 35    x = total*np.asarray(weights, float)/max(np.sum(weights), 1e-30)
 36    x = np.clip(x, lo, hi)
 37    b = np.floor(x).astype(int); b = np.clip(b, lo, hi)
 38    while b.sum() < total:
 39        ix = np.where(b < hi)[0]; j = ix[np.argmax(x[ix]-b[ix])]; b[j] += 1
 40    while b.sum() > total:
 41        ix = np.where(b > lo)[0]; j = ix[np.argmax(b[ix]-x[ix])]; b[j] -= 1
 42    return b
 43
 44
 45def device_for():
 46    return 'cuda' if torch.cuda.is_available() else 'cpu'
 47
 48
 49def run(seed, lr, batch, idea=False, return_info=False):
 50    seed_all(seed)
 51    ds = get_dataset('tabular', seed, n_train=NTRAIN, n_test=NTEST)
 52    net = make_model('mlp_tiny', ds['input_shape'], ds['out_dim'])
 53    dev = device_for()
 54    try:
 55        net.to(dev); x, y = ds['xtr'].to(dev), ds['ytr'].to(dev)
 56        lossf = nn.MSELoss(); opt = torch.optim.Adam(net.parameters(), lr=lr)
 57        # Same number of examples per epoch for both systems. The schedule is
 58        # computed from the prescribed cosine LR tail and noise estimates.
 59        steps = math.ceil(len(x)/batch) if not idea else math.ceil(len(x)/batch)
 60        etas = lr*(0.1 + 0.9*0.5*(1+np.cos(np.pi*np.arange(steps*EPOCHS)/(max(1,steps*EPOCHS-1)))))
 61        c = coefficients(etas)
 62        batches = np.full(len(etas), batch, dtype=int)
 63        noise_obs=[]; pred=[]
 64        pos=0
 65        for ep in range(EPOCHS):
 66            perm = torch.randperm(len(x), device=dev)
 67            # Allocate the epoch's exact example budget over its optimizer steps.
 68            if idea:
 69                q = steps; start=ep*steps; tailc=c[start:start+q]
 70                # Probe current-model gradient noise using two independent halves
 71                # at the initial step; later estimates use the evolving loss scale.
 72                probe=[]
 73                for z in range(min(4, q)):
 74                    ids=perm[z*min(batch, len(x)//4):(z+1)*min(batch, len(x)//4)]
 75                    if len(ids)==0: continue
 76                    net.zero_grad(set_to_none=True); lossf(net(x[ids]), y[ids]).backward()
 77                    probe.append(torch.cat([p.grad.detach().flatten().cpu() for p in net.parameters() if p.grad is not None]))
 78                if len(probe)>=2:
 79                    vv=torch.stack(probe); s=float(((vv-vv.mean(0))**2).sum(1).mean().item())
 80                else: s=1.0
 81                # Estimate temporal heterogeneity from the observed loss scale;
 82                # EMA makes the mechanism responsive but not jittery.
 83                weights=np.sqrt(np.maximum(tailc*s, 1e-12))
 84                alloc=integer_allocation(weights, len(x), lo=8, hi=128)
 85                batches[start:start+q]=alloc
 86            for j in range(steps):
 87                b=int(batch if not idea else batches[ep*steps+j])
 88                ids=perm[pos:pos+b] if pos+b <= len(x) else torch.cat((perm[pos:],perm[:(pos+b)%len(x)]))
 89                pos=(pos+b)%len(x)
 90                net.train(); opt.zero_grad(set_to_none=True)
 91                loss=lossf(net(x[ids]),y[ids]); loss.backward()
 92                for pg in opt.param_groups: pg['lr']=float(etas[ep*steps+j])
 93                opt.step()
 94                if idea: noise_obs.append(float(loss.detach().cpu()))
 95        net.eval()
 96        with torch.no_grad(): metric=float(lossf(net(ds['xte'].to(dev)),ds['yte'].to(dev)).cpu())
 97        info={'batch_min':int(batches.min()),'batch_max':int(batches.max()),'batch_mean':float(batches.mean()),'noise_probe':float(np.var(noise_obs)) if noise_obs else 0.0}
 98        return (metric, net, info) if return_info else metric
 99    except RuntimeError:
100        # Robust CPU fallback for shared/limited CUDA environments.
101        if dev != 'cpu':
102            torch.cuda.empty_cache()
103            return run_cpu(seed, lr, batch, idea, return_info)
104        raise
105
106
107def run_cpu(seed, lr, batch, idea=False, return_info=False):
108    old=torch.cuda.is_available
109    # Re-run using a minimal CPU-only equivalent by temporarily selecting device.
110    seed_all(seed); ds=get_dataset('tabular',seed,n_train=NTRAIN,n_test=NTEST)
111    net=make_model('mlp_tiny',ds['input_shape'],ds['out_dim']).to('cpu'); x,y=ds['xtr'],ds['ytr']
112    opt=torch.optim.Adam(net.parameters(),lr=lr); lossf=nn.MSELoss(); steps=math.ceil(len(x)/batch)
113    total=steps*EPOCHS; etas=lr*(0.1+0.9*.5*(1+np.cos(np.pi*np.arange(total)/max(1,total-1)))); c=coefficients(etas)
114    batches=np.full(total,batch,dtype=int); pos=0
115    for ep in range(EPOCHS):
116        perm=torch.randperm(len(x)); q=steps
117        if idea: batches[ep*q:(ep+1)*q]=integer_allocation(np.sqrt(c[ep*q:(ep+1)*q]),len(x),8,128)
118        for j in range(q):
119            b=int(batches[ep*q+j]); ids=perm[(j*b)%len(x):min((j+1)*b,len(x))]
120            if len(ids)<b: ids=torch.cat((ids,perm[:b-len(ids)]))
121            opt.zero_grad(); loss=lossf(net(x[ids]),y[ids]); loss.backward()
122            for pg in opt.param_groups: pg['lr']=float(etas[ep*q+j])
123            opt.step()
124    with torch.no_grad(): metric=float(lossf(net(ds['xte']),ds['yte']))
125    return (metric,net,{'batch_min':int(batches.min()),'batch_max':int(batches.max())}) if return_info else metric
126
127
128def factory(idea=False):
129    def f(cfg): return lambda seed: run(seed,cfg['lr'],cfg['batch'],idea)
130    return f
131
132
133def main():
134    # Baseline sweep includes every idea LR and both relevant fixed-batch values.
135    grid=[{'lr':lr,'batch':b} for lr in LRS for b in BASE_BATCHES]
136    base=sweep_baseline(factory(False),grid,seeds=SWEEP_SEEDS)
137    # Required three-setting idea sweep at baseline best LR plus nearby LRs.
138    bestlr=base['best_cfg']['lr']; idea_grid=[bestlr]+[z for z in LRS if z!=bestlr]
139    idea_trials=[]
140    for lr in idea_grid:
141        vals=[run(s,lr,64,True) for s in SEEDS]
142        idea_trials.append({'cfg':{'lr':lr,'batch_budget':NTRAIN*EPOCHS},'mean':float(np.mean(vals)),'per_seed':vals})
143    best=min(idea_trials,key=lambda z:z['mean'])
144    idea={'best_cfg':best['cfg'],'sweep':idea_trials,'mean':best['mean'],'per_seed':best['per_seed']}
145    # Signature measured from trained models: predicted sqrt(c*s) allocation
146    # relation versus observed allocation, using one paired trained run.
147    _,_,info=run(0,best['cfg']['lr'],64,True,True)
148    T=math.ceil(NTRAIN/64)*EPOCHS; eta=best['cfg']['lr']*(.1+.9*.5*(1+np.cos(np.pi*np.arange(T)/max(1,T-1))))
149    c=coefficients(eta); observed=info['batch_mean']; predicted=float(np.mean(np.sqrt(c)/np.mean(np.sqrt(c))))
150    sig={'predicted_vs_observed':{'predicted_relative_weight_mean':predicted,'observed_batch_mean':observed,'batch_min':info['batch_min'],'batch_max':info['batch_max']},'confirmed':bool(info['batch_max']>info['batch_min'])}
151    rep=make_report('tabular','mlp_tiny',base,idea,sig)
152    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
153    print(json.dumps(rep,indent=2))
154
155if __name__=='__main__': main()