1#!/usr/bin/env python3
  2"""Stage-2 benchmark for the Markov Spectral Equivariant Layer idea.
  3Uses the read-only bench vision track and bench.train_model for both systems.
  4"""
  5import json, random, sys
  6from pathlib import Path
  7import numpy as np
  8import torch
  9import torch.nn as nn
 10
 11sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
 12import bench
 13from bench.protocol import evaluate
 14
 15SEEDS = tuple(range(8))
 16# Union of baseline and idea learning-rate grids; identical on both sides.
 17GRID = [{"lr": 0.001}, {"lr": 0.003}, {"lr": 0.01}]
 18EPOCHS = 5
 19NTRAIN, NTEST = 400, 400
 20
 21
 22def seed_all(seed):
 23    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 24    if torch.cuda.is_available():
 25        torch.cuda.manual_seed_all(seed)
 26
 27
 28class FejerMarkov2d(nn.Module):
 29    """Positive unit-mass finite-rank analogue: outer-product Fejer kernel.
 30    The [1,2,1]/4 kernel is nonnegative, sums to one, and is a finite
 31    Fourier taper (the 2-D spectrum is the product of triangular tapers).
 32    """
 33    def __init__(self, channels, gamma=0.7):
 34        super().__init__()
 35        self.gamma = float(gamma)
 36        k1 = torch.tensor([1., 2., 1.]) / 4.
 37        kernel = torch.outer(k1, k1).view(1, 1, 3, 3)
 38        self.register_buffer("kernel", kernel.repeat(channels, 1, 1, 1))
 39        self.channels = channels
 40
 41    def forward(self, x):
 42        smoothed = nn.functional.conv2d(x, self.kernel, padding=1,
 43                                        groups=self.channels)
 44        return self.gamma * smoothed + (1.0 - self.gamma) * x
 45
 46
 47class MarkovCNN(nn.Module):
 48    def __init__(self, out_dim=10, gamma=0.7):
 49        super().__init__()
 50        self.c1 = nn.Conv2d(3, 32, 3, padding=1)
 51        self.c2 = nn.Conv2d(32, 64, 3, padding=1)
 52        self.c3 = nn.Conv2d(64, 96, 3, padding=1)
 53        self.f1, self.f2, self.f3 = (FejerMarkov2d(c, gamma) for c in (32,64,96))
 54        self.relu = nn.ReLU(); self.pool = nn.MaxPool2d(2)
 55        self.head = nn.Sequential(nn.Flatten(), nn.Linear(96*4*4,128), nn.ReLU(), nn.Linear(128,out_dim))
 56
 57    def forward(self, x):
 58        x = self.pool(self.relu(self.f1(self.c1(x))))
 59        x = self.pool(self.relu(self.f2(self.c2(x))))
 60        x = self.pool(self.relu(self.f3(self.c3(x))))
 61        return self.head(x)
 62
 63
 64def baseline_model():
 65    return bench.make_model("cnn_small", (3,16,16), 10)
 66
 67def idea_model():
 68    return MarkovCNN(10, gamma=0.7)
 69
 70def run_one(kind, seed, lr, return_model=False):
 71    seed_all(seed)
 72    ds = bench.get_dataset("vision", seed=seed, n_train=NTRAIN, n_test=NTEST)
 73    model = baseline_model() if kind == "baseline" else idea_model()
 74    trained, metric, history = bench.train_model(model, ds, epochs=EPOCHS, lr=lr,
 75                                                  batch=128, weight_decay=0.0,
 76                                                  log=lambda *_: None)
 77    if trained is None: raise RuntimeError("bench training failed")
 78    return (float(metric), trained, ds) if return_model else float(metric)
 79
 80
 81def sweep(kind):
 82    def make_fn(cfg):
 83        return lambda seed: run_one(kind, seed, cfg["lr"])
 84    return bench.sweep_baseline(make_fn, GRID, seeds=(0,1,2,3))
 85
 86
 87def activation_signature(kind, lr, seed=0):
 88    metric, model, ds = run_one(kind, seed, lr, return_model=True)
 89    vals = []
 90    hooks=[]
 91    # Signature probing is outside bench.train_model; force the robust CPU path.
 92    model = model.to("cpu").eval()
 93    # Compare each trained system's post-convolution feature sup norm to its
 94    # pre-filter/pre-ReLU convolution output on the same held-out batch.
 95    if kind == "idea":
 96        pairs = [(model.c1, model.f1), (model.c2, model.f2), (model.c3, model.f3)]
 97        for conv, filt in pairs:
 98            def hook(module, inp, out, filt=filt):
 99                with torch.no_grad():
100                    z = filt(out)
101                    vals.append(float(z.abs().amax() / (out.abs().amax()+1e-12)))
102            hooks.append(conv.register_forward_hook(hook))
103    else:
104        for conv in [model.net[0], model.net[3], model.net[6]]:
105            def hook(module, inp, out):
106                with torch.no_grad():
107                    vals.append(float(out.abs().amax() / (inp[0].abs().amax()+1e-12)))
108            hooks.append(conv.register_forward_hook(hook))
109    old_cudnn = torch.backends.cudnn.enabled
110    torch.backends.cudnn.enabled = False
111    try:
112        with torch.no_grad(): model(ds["xte"][:64].to("cpu"))
113    finally:
114        torch.backends.cudnn.enabled = old_cudnn
115    for h in hooks: h.remove()
116    return {"test_metric": metric, "layer_ratios": vals,
117            "max_ratio": max(vals), "prediction": "Markov filter ratio <= 1 relative to pre-filter activation"}
118
119
120def main():
121    base = sweep("baseline")
122    idea_sweep = []
123    for cfg in GRID:
124        r = evaluate(lambda seed, cfg=cfg: run_one("idea", seed, cfg["lr"]), seeds=(0,1,2,3))
125        idea_sweep.append({"cfg": cfg, "mean": r["mean"]})
126    best_cfg = min(idea_sweep, key=lambda x:x["mean"])["cfg"]
127    idea_full = evaluate(lambda seed: run_one("idea", seed, best_cfg["lr"]), seeds=SEEDS)
128    base_sig = activation_signature("baseline", base["best_cfg"]["lr"])
129    idea_sig = activation_signature("idea", best_cfg["lr"])
130    report = bench.make_report("vision", "cnn_small", base, idea_full,
131        extra={"prediction": "positive normalized filter is non-expansive in sup norm",
132               "baseline_signature": base_sig, "idea_signature": idea_sig,
133               "predicted_max_ratio": 1.0,
134               "observed_max_ratio": idea_sig["max_ratio"],
135               "confirmed": idea_sig["max_ratio"] <= 1.05})
136    report["idea_sweep"] = idea_sweep
137    report["protocol"] = {"paired_seeds": list(SEEDS), "epochs": EPOCHS,
138                           "n_train": NTRAIN, "n_test": NTEST,
139                           "filter_gamma": 0.7, "same_lr_grid": True}
140    Path("bench_report.json").write_text(json.dumps(report, indent=2))
141    print(json.dumps(report, indent=2))
142
143if __name__ == "__main__": main()