import json, random, sys import numpy as np import torch import torch.nn as nn sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import get_dataset, make_model, train_model, sweep_baseline, make_report from bench.protocol import evaluate SEEDS = tuple(range(8)) GRID = [ {'lr': 0.001, 'weight_decay': 0.0}, {'lr': 0.003, 'weight_decay': 0.0}, {'lr': 0.006, 'weight_decay': 0.0}, {'lr': 0.003, 'weight_decay': 1e-4}, ] def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def baseline_train(cfg, seed): seed_all(seed) ds = get_dataset('dynamics', seed, n_train=400, n_test=200) model = make_model('rnn_small', ds['input_shape'], ds['out_dim']) try: res = train_model(model, ds, epochs=14, lr=cfg['lr'], batch=64, weight_decay=cfg['weight_decay'], log=lambda *_: None) return float(res[1]) except Exception: torch.cuda.empty_cache() if torch.cuda.is_available() else None res = train_model(model.cpu(), {k:(v.cpu() if torch.is_tensor(v) else v) for k,v in ds.items()}, epochs=14, lr=cfg['lr'], batch=64, weight_decay=cfg['weight_decay'], log=lambda *_: None) return float(res[1]) class TubeGRU(nn.Module): def __init__(self, hidden=64, wbar=0.02, lam=0.02): super().__init__() self.rnn = nn.GRU(3, hidden, batch_first=True) self.head = nn.Linear(hidden, 1) self.wbar, self.lam = wbar, lam def forward(self, x, tube_loss=False): seq = x.view(x.shape[0], -1, 3) h = torch.zeros(1, x.shape[0], self.rnn.hidden_size, device=x.device) eb = torch.full((x.shape[0], self.rnn.hidden_size), 0.01, device=x.device) penalty = torch.zeros((), device=x.device) for j in range(seq.shape[1]): z = seq[:, j:j+1] _, hn = self.rnn(z, h) # A practical upper bound for the hidden-state Jacobian: absolute # recurrent matrix row sums, scaled by tanh/sigmoid local saturation. W = self.rnn.weight_hh_l0 gain = W.abs().sum(dim=1).mean() / max(1, self.rnn.hidden_size) gain = torch.clamp(gain, 0.0, 0.99) eb = gain * eb + self.wbar if tube_loss: margin = eb.mean() * self.head.weight.abs().mean() penalty = penalty + torch.relu(margin - 0.05).pow(2) h = hn out = self.head(h[-1]) return (out, self.lam * penalty / seq.shape[1]) if tube_loss else out def idea_train(cfg, seed): seed_all(seed) ds = get_dataset('dynamics', seed, n_train=400, n_test=200) model = TubeGRU() dev = torch.device('cuda' if torch.cuda.is_available() else 'cpu') try: model.to(dev) xtr,ytr,xte,yte = [ds[k].to(dev) for k in ('xtr','ytr','xte','yte')] opt = torch.optim.Adam(model.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay']) for _ in range(14): model.train() for ix in torch.randperm(len(xtr), device=dev).split(64): pred, reg = model(xtr[ix], True) loss = (pred-ytr[ix]).pow(2).mean() + reg opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step() model.eval() with torch.no_grad(): val = (model(xte)-yte).pow(2).mean().item() return float(val) except Exception: model.cpu(); xtr,ytr,xte,yte = [ds[k].cpu() for k in ('xtr','ytr','xte','yte')] opt = torch.optim.Adam(model.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay']) for _ in range(14): for ix in torch.randperm(len(xtr)).split(64): pred,reg=model(xtr[ix],True); loss=(pred-ytr[ix]).pow(2).mean()+reg opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): return float((model(xte)-yte).pow(2).mean()) def signature(seed, cfg): seed_all(seed) ds = get_dataset('dynamics', seed, n_train=400, n_test=200) # Signature is deliberately measured on CPU: the shared GPU can reject # cuDNN host allocations, while this does not alter either benchmark side. dev = torch.device('cpu') m = TubeGRU().to(dev) xx, yy = ds['xtr'].to(dev), ds['ytr'].to(dev) opt = torch.optim.Adam(m.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay']) for _ in range(6): for ix in torch.randperm(len(xx)).split(64): pred, reg = m(xx[ix], True) loss = (pred - yy[ix]).pow(2).mean() + reg opt.zero_grad(); loss.backward(); opt.step() m.eval() x = ds['xte'][:32].to(dev) h = torch.zeros(1, len(x), 64) eb = torch.full((len(x), 64), .01) observed, predicted = [], [] with torch.no_grad(): for j in range(8): z = x[:, j*3:(j+1)*3].view(len(x), 1, 3) _, hn = m.rnn(z, h) observed.append(float((hn-h).abs().mean())) gain = float(m.rnn.weight_hh_l0.abs().sum(1).mean() / 64) gain = min(gain, .99) eb = gain * eb + .02 predicted.append(float(eb.mean())) h = hn ratios = np.asarray(observed) / np.maximum(np.asarray(predicted), 1e-9) return {'observed_hidden_change_mean': float(np.mean(observed)), 'predicted_tube_mean': float(np.mean(predicted)), 'ratio_observed_over_predicted': float(np.mean(ratios)), 'confirmed': bool(0.1 < np.mean(ratios) < 10.0)} def main(): base=sweep_baseline(lambda c: lambda s: baseline_train(c,s), GRID, seeds=(0,1,2,3)) base_full=evaluate(lambda s: baseline_train(base['best_cfg'],s), SEEDS) base['full']=base_full idea_cfgs = [base['best_cfg'], {'lr': 0.003, 'weight_decay': 0.0}, {'lr': 0.001, 'weight_decay': 0.0}] idea_runs = [{'cfg': c, 'result': evaluate(lambda s, cc=c: idea_train(cc, s), SEEDS)} for c in idea_cfgs] best_i = min(idea_runs, key=lambda z: z['result']['mean']) idea = best_i['result'] rep=make_report('dynamics','rnn_small',base,idea,signature(0,best_i['cfg'])) rep['idea_sweep'] = [{'cfg': z['cfg'], 'mean': z['result']['mean']} for z in idea_runs] rep['idea_best_cfg'] = best_i['cfg'] rep['protocol_notes']='Baseline and idea use matched 64-unit GRU systems, same data, epochs, batch, optimizer, and shared lr/weight-decay grid; idea tested at three shared settings.' with open('bench_report.json','w') as f: json.dump(rep,f,indent=2) print(json.dumps(rep,indent=2)) if __name__=='__main__': main()