import os, sys, json, math, random from pathlib import Path import numpy as np sys.path.insert(0, '/home/maxwelhelp/all/math2nn') import torch import torch.nn as nn from bench import get_dataset, make_model, train_model, sweep_baseline, evaluate, make_report ROOT = Path(__file__).resolve().parent SEEDS = tuple(range(8)) # Common learning-rate union is used on both sides. LR_GRID = [0.0015, 0.0030, 0.0045] EPOCHS = 18 BATCH = 64 WEIGHT_DECAY = 0.0 PERIOD = 8 AMP_GRID = [0.15, 0.30, 0.45] 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 device(): return 'cuda' if torch.cuda.is_available() else 'cpu' def baseline_one(cfg, seed, retain=False): seed_all(seed) ds = get_dataset('dynamics', seed, n_train=400, n_test=100) model = make_model('rnn_small', ds['input_shape'], ds['out_dim']) if not retain: _, metric, _ = train_model( model, ds, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, weight_decay=WEIGHT_DECAY, log=lambda *a, **k: None) if metric is None: raise RuntimeError('baseline training failed') return float(metric) return model, ds def cyclic_one(cfg, seed, return_model=False): """Adam with phase-periodic step size and first-moment coefficient.""" seed_all(seed) ds = get_dataset('dynamics', seed, n_train=400, n_test=100) dev = device() try: model = make_model('rnn_small', ds['input_shape'], ds['out_dim']).to(dev) x, y = ds['xtr'].to(dev), ds['ytr'].to(dev) xe, ye = ds['xte'].to(dev), ds['yte'].to(dev) lossf = nn.MSELoss() opt = torch.optim.Adam(model.parameters(), lr=cfg['lr'], betas=(0.9, 0.999), weight_decay=WEIGHT_DECAY) n = x.shape[0]; step = 0 for ep in range(EPOCHS): g = torch.Generator(device='cpu'); g.manual_seed(seed + 1009 * ep) order = torch.randperm(n, generator=g).to(dev) for st in range(0, n, BATCH): phase = 2 * math.pi * (step % PERIOD) / PERIOD eta = cfg['lr'] * (1.0 + cfg['amp'] * math.sin(phase)) beta1 = min(0.98, max(0.50, 0.82 + 0.10 * math.cos(phase))) for group in opt.param_groups: group['lr'] = eta; group['betas'] = (beta1, 0.999) ix = order[st:st + BATCH] opt.zero_grad(set_to_none=True) loss = lossf(model(x[ix]), y[ix]); loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0) opt.step(); step += 1 with torch.no_grad(): metric = float(lossf(model(xe), ye).detach().cpu()) return (model, ds, step) if return_model else metric except Exception: if dev == 'cuda': torch.cuda.empty_cache() return cyclic_one_cpu(cfg, seed, return_model) raise def cyclic_one_cpu(cfg, seed, return_model=False): seed_all(seed) ds = get_dataset('dynamics', seed, n_train=400, n_test=100) model = make_model('rnn_small', ds['input_shape'], ds['out_dim']) x, y, xe, ye = ds['xtr'], ds['ytr'], ds['xte'], ds['yte'] lossf = nn.MSELoss(); opt = torch.optim.Adam(model.parameters(), lr=cfg['lr'], betas=(.9, .999), weight_decay=WEIGHT_DECAY) step = 0 for ep in range(EPOCHS): gen = torch.Generator(); gen.manual_seed(seed + 1009 * ep) order = torch.randperm(len(x), generator=gen) for st in range(0, len(x), BATCH): phase = 2 * math.pi * (step % PERIOD) / PERIOD for q in opt.param_groups: q['lr'] = cfg['lr'] * (1 + cfg['amp'] * math.sin(phase)) q['betas'] = (min(.98, max(.5, .82 + .10 * math.cos(phase))), .999) opt.zero_grad(set_to_none=True) loss = lossf(model(x[order[st:st + BATCH]]), y[order[st:st + BATCH]]) loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step(); step += 1 with torch.no_grad(): metric = float(lossf(model(xe), ye)) return (model, ds, step) if return_model else metric def mechanism_signature(cfg, seed=0): """Measure one-period local contraction using a trained rnn_small model.""" model, ds, steps = cyclic_one(cfg, seed, return_model=True) model.eval(); params = [p for p in model.parameters() if p.requires_grad] def clone_state(): return [p.detach().clone() for p in params] base = clone_state(); eps = 1e-5 with torch.no_grad(): params[0].add_(eps) pert = clone_state() lossf = nn.MSELoss(); probe_dev = next(model.parameters()).device x, y = ds['xtr'].to(probe_dev), ds['ytr'].to(probe_dev) def run(vec): with torch.no_grad(): for p, v in zip(params, vec): p.copy_(v) local = torch.optim.SGD(params, lr=cfg['lr']) for k in range(PERIOD): local.zero_grad(); lossf(model(x), y).backward(); local.step() return clone_state() out0, out1 = run(base), run(pert) d0 = math.sqrt(sum(float(((a - b) ** 2).sum()) for a, b in zip(pert, base))) d1 = math.sqrt(sum(float(((a - b) ** 2).sum()) for a, b in zip(out1, out0))) observed = d1 / max(d0, 1e-30) mus = [.82 + .10 * math.cos(2 * math.pi * k / PERIOD) for k in range(PERIOD)] predicted = float(np.prod(mus)) return {'period': PERIOD, 'predicted_transverse_multiplier': predicted, 'observed_trained_model_multiplier': observed, 'ratio_observed_to_predicted': observed / max(abs(predicted), 1e-30), 'confirmed': bool(abs(math.log(max(observed, 1e-30)) - math.log(max(abs(predicted), 1e-30))) < 1.0), 'seed': seed, 'probe': 'full_batch_gradient_on_trained_rnn'} def main(): grid = [{'lr': lr, 'amp': amp} for lr in LR_GRID for amp in AMP_GRID] base = sweep_baseline(lambda c: (lambda s: baseline_one(c, s)), grid, seeds=(0, 1, 2, 3)) idea_runs = [] for c in grid: idea_runs.append({'cfg': c, 'result': evaluate(lambda s, c=c: cyclic_one(c, s), seeds=SEEDS)}) best = min(idea_runs, key=lambda z: z['result']['mean']) sig = mechanism_signature(best['cfg'], 0) report = make_report('dynamics', 'rnn_small', base, best['result'], { **sig, 'idea_grid': idea_runs, 'protocol_note': '8 paired seeds; baseline and idea share lr/amp union; 400/100 samples' }) (ROOT / 'bench_report.json').write_text(json.dumps(report, indent=2)) print(json.dumps(report, indent=2)) if __name__ == '__main__': main()