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, evaluate, sweep_baseline, make_report
  9
 10OUT = Path(__file__).with_name('bench_report.json')
 11EPOCHS = 18
 12BATCH = 128
 13RHO_AUG = 0.50
 14DENSITIES = [0.25, 0.5, 0.75]  # baseline real-space density; union with idea lr grid is shared
 15LRS = [1e-3, 3e-3, 1e-2]
 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
 23
 24def device():
 25    return 'cuda' if torch.cuda.is_available() else 'cpu'
 26
 27
 28def flatten_params(net):
 29    # Prune every affine weight tensor, excluding biases. Keep tensor shapes for forward.
 30    refs = []
 31    for mod in net.modules():
 32        if isinstance(mod, nn.Linear): refs.append(mod)
 33    return refs
 34
 35
 36def run_one(seed, lr, augmented, density=0.5):
 37    seed_all(seed)
 38    ds = get_dataset('tabular', seed=seed, n_train=400, n_test=200)
 39    # Make identical random base architecture/initialization for paired systems.
 40    net = make_model('mlp_tiny', ds['input_shape'], ds['out_dim'])
 41    refs = flatten_params(net)
 42    dev = device()
 43    try:
 44        net = net.to(dev)
 45        xtr = torch.as_tensor(ds['xtr'], dtype=torch.float32, device=dev)
 46        ytr = torch.as_tensor(ds['ytr'], dtype=torch.float32, device=dev)
 47        xte = torch.as_tensor(ds['xte'], dtype=torch.float32, device=dev)
 48        yte = torch.as_tensor(ds['yte'], dtype=torch.float32, device=dev)
 49        # Freeze all weights; score parameters are the only trainable quantities.
 50        for p in net.parameters(): p.requires_grad_(False)
 51        real_scores = [nn.Parameter(torch.randn_like(m.weight, device=dev)) for m in refs]
 52        dummy_scores = [nn.Parameter(torch.randn_like(m.weight, device=dev)) for m in refs] if augmented else []
 53        params = real_scores + dummy_scores
 54        opt = torch.optim.Adam(params, lr=lr)
 55        # Exact fixed candidate quota. For a doubled layer, K=floor(rho_aug*2M).
 56        quotas = []
 57        for m in refs:
 58            M = m.weight.numel()
 59            quotas.append(max(1, min(2*M if augmented else M,
 60                                     int(math.floor((RHO_AUG if augmented else density) * (2*M if augmented else M))))))
 61        for epoch in range(EPOCHS):
 62            perm = torch.randperm(xtr.shape[0], device=dev)
 63            for start in range(0, xtr.shape[0], BATCH):
 64                ix = perm[start:start+BATCH]
 65                hlist = []
 66                for m, sr, sd, K in zip(refs, real_scores, dummy_scores if augmented else [None]*len(refs), quotas):
 67                    sa = torch.cat([sr.reshape(-1), sd.reshape(-1)]) if augmented else sr.reshape(-1)
 68                    hard = torch.zeros_like(sa); hard[torch.topk(sa, K).indices] = 1.
 69                    st = (hard - sa).detach() + sa
 70                    hlist.append(st[:m.weight.numel()].reshape_as(m.weight))
 71                # Functional forward with masked weights, preserving the shared MLP.
 72                z = xtr[ix]
 73                li = 0
 74                for mod in net.modules():
 75                    if isinstance(mod, nn.Linear):
 76                        z = torch.nn.functional.linear(z, mod.weight * hlist[li], mod.bias)
 77                        li += 1
 78                        if li < len(refs): z = torch.relu(z)
 79                loss = torch.mean((z - ytr[ix]) ** 2)
 80                opt.zero_grad(); loss.backward(); opt.step()
 81        with torch.no_grad():
 82            masks=[]; real_counts=[]; dummy_counts=[]
 83            for m, sr, sd, K in zip(refs, real_scores, dummy_scores if augmented else [None]*len(refs), quotas):
 84                sa = torch.cat([sr.reshape(-1), sd.reshape(-1)]) if augmented else sr.reshape(-1)
 85                inds = torch.topk(sa, K).indices
 86                mask = torch.zeros(m.weight.numel(), device=inds.device); real_ind = inds[inds < m.weight.numel()]; mask[real_ind] = 1.; masks.append(mask.reshape_as(m.weight))
 87                real_counts.append(int((inds < m.weight.numel()).sum().item()))
 88                dummy_counts.append(int(K - real_counts[-1]))
 89            z=xte; li=0
 90            for mod in net.modules():
 91                if isinstance(mod, nn.Linear):
 92                    z=torch.nn.functional.linear(z, mod.weight*masks[li], mod.bias); li+=1
 93                    if li<len(refs): z=torch.relu(z)
 94            metric=float(torch.mean((z-yte)**2).item())
 95        return metric, {'real_density': float(np.mean([r/m.weight.numel() for r,m in zip(real_counts,refs)])),
 96                        'dummy_fraction': float(sum(dummy_counts)/sum(quotas)), 'real_counts': real_counts}
 97    except Exception as e:
 98        if dev == 'cuda':
 99            torch.cuda.empty_cache()
100            # deterministic CPU retry
101            return run_one_cpu(seed, lr, augmented, density)
102        raise
103
104
105def run_one_cpu(seed, lr, augmented, density=0.5):
106    old=torch.cuda.is_available
107    # Explicit CPU implementation by temporarily using a compact equivalent call is avoided;
108    # rerun after forcing tensors/model to CPU through the shared function's code path.
109    # In normal bench environments CUDA or CPU path above succeeds; this fallback is a direct small model.
110    seed_all(seed); ds=get_dataset('tabular', seed=seed, n_train=400, n_test=200)
111    net=make_model('mlp_tiny', ds['input_shape'], ds['out_dim']); refs=flatten_params(net)
112    xtr=torch.tensor(ds['xtr'],dtype=torch.float32); ytr=torch.tensor(ds['ytr'],dtype=torch.float32)
113    xte=torch.tensor(ds['xte'],dtype=torch.float32); yte=torch.tensor(ds['yte'],dtype=torch.float32)
114    for p in net.parameters(): p.requires_grad_(False)
115    sr=[nn.Parameter(torch.randn_like(m.weight)) for m in refs]; sd=[nn.Parameter(torch.randn_like(m.weight)) for m in refs] if augmented else []
116    opt=torch.optim.Adam(sr+sd,lr=lr); qs=[max(1,int(math.floor((RHO_AUG if augmented else density)*(2*m.weight.numel() if augmented else m.weight.numel())))) for m in refs]
117    for _ in range(EPOCHS):
118      for st in range(0,len(xtr),BATCH):
119        hs=[]
120        for m,a,b,k in zip(refs,sr,sd if augmented else [None]*len(refs),qs):
121          q=torch.cat([a.flatten(),b.flatten()]) if augmented else a.flatten(); hard=torch.zeros_like(q); hard[torch.topk(q,k).indices]=1; hs.append(((hard-q).detach()+q)[:m.weight.numel()].reshape_as(m.weight))
122        z=xtr[st:st+BATCH]; j=0
123        for mod in net.modules():
124          if isinstance(mod,nn.Linear): z=nn.functional.linear(z,mod.weight*hs[j],mod.bias); j+=1; z=torch.relu(z) if j<len(refs) else z
125        loss=((z-ytr[st:st+BATCH])**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
126    with torch.no_grad():
127      hs=[]; rc=[]; dc=[]
128      for m,a,b,k in zip(refs,sr,sd if augmented else [None]*len(refs),qs):
129        q=torch.cat([a.flatten(),b.flatten()]) if augmented else a.flatten(); ind=torch.topk(q,k).indices; mask=torch.zeros(m.weight.numel()); real_ind=ind[ind<m.weight.numel()]; mask[real_ind]=1.; hs.append(mask.reshape_as(m.weight)); rc.append(int(real_ind.numel())); dc.append(k-rc[-1])
130      z=xte;j=0
131      for mod in net.modules():
132       if isinstance(mod,nn.Linear): z=nn.functional.linear(z,mod.weight*hs[j],mod.bias);j+=1;z=nn.functional.relu(z) if j<len(refs) else z
133      return float(((z-yte)**2).mean()), {'real_density':float(np.mean([r/m.weight.numel() for r,m in zip(rc,refs)])),'dummy_fraction':float(sum(dc)/sum(qs)),'real_counts':rc}
134
135
136def main():
137    # Baseline method knobs (density and lr) and all idea lrs share the same union.
138    grid=[{'lr':lr,'density':rho} for lr in LRS for rho in DENSITIES]
139    def mk(cfg): return lambda s: run_one(s,cfg['lr'],False,cfg['density'])[0]
140    base=sweep_baseline(mk,grid)
141    # Three idea settings: baseline-best lr plus two nearby values, all in baseline grid.
142    idea_lrs=sorted(set(LRS))
143    idea_trials=[]
144    for lr in idea_lrs:
145        r=evaluate(lambda s, lr=lr: run_one(s,lr,True)[0], seeds=(0,1,2,3))
146        idea_trials.append({'cfg':{'lr':lr,'rho_aug':RHO_AUG},'mean':r['mean']})
147    best_idea_lr=min(idea_trials,key=lambda x:x['mean'])['cfg']['lr']
148    idea=evaluate(lambda s: run_one(s,best_idea_lr,True)[0])
149    sig=[run_one(s,best_idea_lr,True)[1] for s in range(8)]
150    sig_b=[run_one(s,base['best_cfg']['lr'],False,base['best_cfg']['density'])[1] for s in range(8)]
151    obs_dummy=float(np.mean([x['dummy_fraction'] for x in sig]))
152    obs_density=float(np.mean([x['real_density'] for x in sig]))
153    signature={'prediction':'For iid initialized real/dummy scores and fixed Top-K, expected dummy fraction is 0.5 and expected real density is rho_aug=0.5.',
154      'predicted_dummy_fraction':0.5,'observed_dummy_fraction_mean':obs_dummy,
155      'predicted_real_density':RHO_AUG,'observed_real_density_mean':obs_density,
156      'baseline_observed_real_density':float(np.mean([x['real_density'] for x in sig_b])),
157      'confirmed':bool(abs(obs_dummy-.5)<.12 and abs(obs_density-RHO_AUG)<.12)}
158    rep=make_report('tabular','mlp_tiny',base,idea,{'track_choice':'tabular matches pruning/regularization of a frozen MLP; identical MLP and task are used for both systems.','shared_lr_grid':LRS,'baseline_density_grid':DENSITIES,'idea_sweep':idea_trials,'idea_best_cfg':{'lr':best_idea_lr,'rho_aug':RHO_AUG},'signature':signature})
159    rep['bench_report']=rep.copy()
160    rep['idea']['sweep']=idea_trials
161    OUT.write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
162
163if __name__=='__main__': main()