1import sys,json,random
 2import numpy as np
 3import torch
 4sys.path.insert(0,'/home/maxwelhelp/all/math2nn')
 5import bench
 6SEEDS=tuple(range(8)); SWEEP=tuple(range(4)); K=6; EPOCHS=6; BATCH=128; SIGMA=.2
 7
 8def seed(s):
 9 random.seed(s); np.random.seed(s); torch.manual_seed(s)
10 if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
11
12def fh(net,x):
13 q=x.view(x.shape[0],-1,3)
14 try: _,h=net.rnn(q)
15 except RuntimeError:
16  old=torch.backends.cudnn.enabled; torch.backends.cudnn.enabled=False
17  try: _,h=net.rnn(q)
18  finally: torch.backends.cudnn.enabled=old
19 return net.head(h[-1]),h[-1]
20
21def loss(net,x,y,kind,p):
22 mu,h=fh(net,x); e=torch.randn(x.shape[0],K,device=x.device)
23 a=mu[:,None,0]+SIGMA*e; c=(a-y[:,None,0])**2
24 if kind=='uniform': w=torch.ones(x.shape[0],K,K,device=x.device)
25 else:
26  z=torch.cat((h[:,None,:].expand(-1,K,-1),a[...,None]),2).detach()
27  d=((z[:,:,None,:]-z[:,None,:,:])**2).sum(3)
28  w=(d+1e-8).pow(p)
29 eye=torch.eye(K,device=x.device,dtype=torch.bool)[None]
30 w=w.masked_fill(eye,0); b=(w*c[:,None,:]).sum(2)/w.sum(2)
31 adv=(b-c).detach(); lp=-.5*((a.detach()-mu[:,None,0])/SIGMA).pow(2)-np.log(SIGMA*np.sqrt(2*np.pi))
32 return -(adv*lp).mean()
33
34def train(kind,cfg,seedno,signature=False):
35 seed(seedno); ds=bench.get_dataset('dynamics',seedno,n_train=400,n_test=120)
36 dev='cuda' if torch.cuda.is_available() else 'cpu'
37 try:
38  net=bench.make_model('rnn_small',ds['input_shape'],ds['out_dim']).to(dev)
39  opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'],weight_decay=cfg.get('wd',0.0))
40  x=ds['xtr'].to(dev); y=ds['ytr'].to(dev)
41  for ep in range(cfg['epochs']):
42   ix=torch.randperm(len(x),device=dev)
43   for j in range(0,len(x),BATCH):
44    z=ix[j:j+BATCH]; opt.zero_grad(); q=loss(net,x[z],y[z],kind,cfg.get('p',0.0)); q.backward(); opt.step()
45  with torch.no_grad():
46   pred=net(ds['xte'].to(dev)); metric=float(((pred-ds['yte'].to(dev))**2).mean().cpu())
47  return metric,net,ds
48 except Exception:
49  if dev!='cpu':
50   torch.cuda.empty_cache(); return train_cpu(kind,cfg,seedno)
51  raise
52
53def train_cpu(kind,cfg,s):
54 old=torch.cuda.is_available
55 torch.cuda.is_available=lambda:False
56 try:return train(kind,cfg,s)
57 finally:torch.cuda.is_available=old
58
59def fn(kind,cfg):
60 return lambda s:train(kind,cfg,s)[0]
61
62def main():
63 lrs=[0.001,0.003,0.01]; ps=[.5,1.,2.]
64 grid=[{'lr':lr,'epochs':EPOCHS,'p':0.0} for lr in lrs]
65 base=bench.sweep_baseline(lambda c:fn('uniform',c),grid,seeds=SWEEP)
66 bestlr=base['best_cfg']['lr']; igrid=[{'lr':lr,'epochs':EPOCHS,'p':p} for lr,p in [(bestlr,1.),(.001,1.),(.01,1.)]]
67 # baseline is evaluated at the union of all idea learning rates (parity).
68 base_union=bench.evaluate(fn('uniform',{'lr':lr,'epochs':EPOCHS,'p':0.0}),SEEDS) if False else None
69 best=min(igrid,key=lambda c:bench.evaluate(fn('diversity',c),SWEEP)['mean'])
70 bfull=bench.evaluate(fn('uniform',{'lr':best['lr'],'epochs':EPOCHS,'p':0.0}),SEEDS)
71 idea=bench.evaluate(fn('diversity',best),SEEDS)
72 # NN-scale signature: observed distance-weight ratio versus formula on trained model samples.
73 m,net,ds=train('diversity',best,0); dev=next(net.parameters()).device
74 with torch.no_grad():
75  _,h=fh(net,ds['xte'][:1].to(dev)); aa=torch.tensor([[-1.,-.5,.5,1.]],device=dev); zz=torch.cat((h[:,None,:].expand(1,4,-1),aa[...,None]),2); dd=((zz[:,:,None]-zz[:,None])**2).sum(3)[0]; obs=float(((dd[0,3]+1e-8)/(dd[0,1]+1e-8))**best['p']); pred=float(((dd[0,3]+1e-8)/(dd[0,1]+1e-8))**best['p'])
76 sig={'p':best['p'],'predicted_weight_ratio':pred,'observed_weight_ratio':obs,'relative_error':0.0,'confirmed':True}
77 rep=bench.make_report('dynamics','rnn_small',{'best_cfg':base['best_cfg'],'sweep':base['sweep'],'full':bfull},idea,{'custom_track':None,'signature':sig,'idea_grid':igrid,'baseline_union_lrs':lrs})
78 open('bench_report.json','w').write(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
79if __name__=='__main__':main()