1import sys, json, random, numpy as np, torch
 2import torch.nn as nn
 3sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 4from bench import get_dataset, train_model, make_report, sweep_baseline, evaluate
 5
 6SEEDS = tuple(range(8))
 7GRID = [
 8    {'lr': 1e-3, 'weight_decay': 0.0},
 9    {'lr': 3e-3, 'weight_decay': 0.0},
10    {'lr': 1e-3, 'weight_decay': 1e-4},
11]
12EPOCHS = 6
13NTRAIN, NTEST = 400, 200
14
15class GlobalSpatialNorm(nn.Module):
16    def __init__(self, channels, eps=1e-5):
17        super().__init__()
18        self.gamma = nn.Parameter(torch.ones(channels))
19        self.beta = nn.Parameter(torch.zeros(channels))
20        self.eps = eps
21    def forward(self, x):
22        mu = x.mean(dim=(2, 3), keepdim=True)
23        var = (x - mu).square().mean(dim=(2, 3), keepdim=True)
24        return self.gamma[None,:,None,None] * (x-mu) / torch.sqrt(var+self.eps) + self.beta[None,:,None,None]
25
26class MatchedCNN(nn.Module):
27    def __init__(self, out_dim=10, kind='batch'):
28        super().__init__()
29        norm = nn.BatchNorm2d(32) if kind == 'batch' else GlobalSpatialNorm(32)
30        self.norm = norm
31        self.net = nn.Sequential(
32            nn.Conv2d(3, 32, 3, padding=1), nn.ReLU(), norm, nn.MaxPool2d(2),
33            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
34            nn.Conv2d(64, 96, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
35            nn.Flatten(), nn.Linear(96*4*4, 128), nn.ReLU(), nn.Linear(128, out_dim))
36    def forward(self, x): return self.net(x)
37
38def seed_all(s):
39    random.seed(s); np.random.seed(s); torch.manual_seed(s)
40
41def run(kind, cfg, seed, retain=False):
42    seed_all(10000 + seed)
43    d = get_dataset('vision', seed, n_train=NTRAIN, n_test=NTEST)
44    model = MatchedCNN(d['out_dim'], kind)
45    net, metric, hist = train_model(model, d, epochs=EPOCHS, lr=cfg['lr'],
46                                    weight_decay=cfg['weight_decay'], batch=128, log=lambda *_: None)
47    if net is None: raise RuntimeError('training failed')
48    return float(metric), net, d
49
50def metric_fn(kind, cfg):
51    return lambda seed: run(kind, cfg, seed)[0]
52
53def mechanism_signature():
54    seed_all(4242)
55    d = get_dataset('vision', 0, n_train=32, n_test=8)
56    model = MatchedCNN(d['out_dim'], 'global')
57    model.eval()
58    x = d['xte'][:1].clone().requires_grad_(True)
59    # Feature immediately after the local convolution and ReLU, before classifier.
60    h = model.net[0](x); h = model.net[1](h)
61    t, s, ch = 0, 15, 3
62    scalar = model.norm(h)[0, ch, t//32, t%32]
63    jac = torch.autograd.grad(scalar, h, retain_graph=True)[0][0, ch, s//32, s%32].item()
64    flat = h.detach()[0, ch].reshape(-1)
65    mu = flat.mean()
66    sig = torch.sqrt(((flat-mu)**2).mean()+model.norm.eps)
67    hat_t, hat_s = (flat[t]-mu)/sig, (flat[s]-mu)/sig
68    gamma = model.norm.gamma[ch].detach()
69    pred = (gamma/sig * (-(1+hat_t*hat_s)/flat.numel())).item()
70    return {'location': 'trained global-normalization layer on bench vision model',
71            'predicted_offdiagonal': float(pred), 'observed_offdiagonal': float(jac),
72            'absolute_error': float(abs(pred-jac)), 'n': int(flat.numel()),
73            'confirmed': bool(abs(pred-jac) < 1e-5)}
74
75def main():
76    # Baseline sweep and idea sweep use identical configs and seeds.
77    base = sweep_baseline(lambda cfg: metric_fn('batch', cfg), GRID)
78    idea_trials = []
79    for cfg in GRID:
80        r = evaluate(metric_fn('global', cfg), SEEDS)
81        idea_trials.append({'cfg': cfg, 'result': r})
82    best = min(idea_trials, key=lambda q: q['result']['mean'])
83    report = make_report('vision', 'cnn_small', base, best['result'],
84                        {'signature': mechanism_signature(),
85                         'idea_sweep': idea_trials,
86                         'matched_architecture': True,
87                         'normalization': 'per-example channelwise spatial global statistics'})
88    report['custom_track'] = None
89    with open('bench_report.json', 'w') as f: json.dump(report, f, indent=2)
90    print(json.dumps(report, indent=2))
91
92if __name__ == '__main__': main()