1import json, math, random, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6import torch.nn.functional as F
  7
  8SEED = 1729
  9
 10def seed_all(seed=SEED):
 11    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 12    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 13
 14
 15def gaussian_kernel(sigma, device, dtype, truncate=3.0):
 16    radius = max(1, int(math.ceil(truncate*sigma)))
 17    x = torch.arange(-radius, radius+1, device=device, dtype=dtype)
 18    g = torch.exp(-0.5*(x/sigma)**2); g = g/g.sum()
 19    k = g[:, None] * g[None, :]
 20    return k
 21
 22
 23def smooth_depthwise(x, sigma):
 24    b,c,h,w = x.shape
 25    k = gaussian_kernel(float(sigma), x.device, x.dtype)
 26    p = k.shape[0]//2
 27    weight = k[None,None].expand(c,1,-1,-1)
 28    return F.conv2d(F.pad(x, (p,p,p,p), mode='replicate'), weight, groups=c)
 29
 30
 31def besov_regularizer(x, sigma0=1.0, r=2.0, J=3, s=0.5, return_scales=False):
 32    scales = [smooth_depthwise(x, sigma0*(r**j)) for j in range(J+1)]
 33    ds = [scales[j]-scales[j+1] for j in range(J)]
 34    # RMS over spatial dimensions, then average channels and batch for a scalar.
 35    vals = []
 36    for j,d in enumerate(ds):
 37        rms = d.pow(2).mean(dim=(-2,-1)).sqrt() # B,C
 38        vals.append((2.0**(j*s))*rms)
 39    stacked = torch.stack(vals, dim=0)
 40    reg = torch.sqrt((stacked**2).sum(dim=0)+1e-12).mean()
 41    if return_scales: return reg, scales, ds, stacked
 42    return reg
 43
 44
 45def math_check():
 46    seed_all(); x = torch.randn(2,3,32,32)
 47    reg, scales, ds, weighted = besov_regularizer(x, return_scales=True)
 48    residual = (scales[0]-scales[-1]-sum(ds)).abs().max().item()
 49    # A sinusoid's response should move toward earlier (fine) scales as frequency rises.
 50    yy,xx = torch.meshgrid(torch.arange(32), torch.arange(32), indexing='ij')
 51    responses=[]
 52    for cycles in [1,4,10]:
 53        z = torch.sin(2*math.pi*cycles*xx/32).float()[None,None]
 54        _,_,_,w = besov_regularizer(z, return_scales=True)
 55        responses.append(w.detach().numpy().ravel().tolist())
 56    # Compare s=0 versus positive s on the same nonzero multi-scale signal.
 57    low = torch.sin(2*math.pi*2*xx/32).float()[None,None]
 58    high = torch.sin(2*math.pi*12*xx/32).float()[None,None]
 59    r0_low=besov_regularizer(low,s=0.0).item(); r0_high=besov_regularizer(high,s=0.0).item()
 60    rp_low=besov_regularizer(low,s=.5).item(); rp_high=besov_regularizer(high,s=.5).item()
 61    return {'telescoping_max_abs':residual, 'sinusoid_weighted_by_scale':responses,
 62            's0_low_high':[r0_low,r0_high], 's05_low_high':[rp_low,rp_high]}
 63
 64class TinyNet(nn.Module):
 65    def __init__(self, use_reg=False, s=.5, lam=0.01):
 66        super().__init__(); self.use_reg=use_reg; self.s=s; self.lam=lam
 67        self.c1=nn.Conv2d(1,16,3,padding=1); self.c2=nn.Conv2d(16,32,3,padding=1)
 68        self.head=nn.Linear(32*8*8,2)
 69    def forward(self,x):
 70        x=F.relu(self.c1(x)); x=F.avg_pool2d(x,2)
 71        feat=F.relu(self.c2(x)); x=F.avg_pool2d(feat,2)
 72        logits=self.head(x.flatten(1))
 73        reg=besov_regularizer(feat,s=self.s) if self.use_reg else feat.new_zeros(())
 74        return logits,reg
 75
 76def dataset(n, seed):
 77    g=torch.Generator().manual_seed(seed)
 78    # Two classes: a persistent broad bar versus a fine checker/stripe texture,
 79    # with nuisance translations and Gaussian noise.
 80    imgs=torch.randn(n,1,32,32,generator=g)*.12; ys=torch.randint(0,2,(n,),generator=g)
 81    for i,y in enumerate(ys.tolist()):
 82        shift=int(torch.randint(-2,3,(1,),generator=g));
 83        if y==0:
 84            imgs[i,:,9+shift:23+shift,14:18] += 0.9
 85            imgs[i,:,14:18,9+shift:23+shift] += 0.9
 86        else:
 87            yy,xx=torch.meshgrid(torch.arange(32),torch.arange(32),indexing='ij')
 88            imgs[i,0] += .38*torch.sin(2*math.pi*7*xx/32)
 89            imgs[i,0] += .38*torch.sin(2*math.pi*7*yy/32)
 90    return imgs,ys
 91
 92def run_train(use_reg, lam, s, train, test, device):
 93    seed_all(99); model=TinyNet(use_reg,s,lam).to(device)
 94    opt=torch.optim.AdamW(model.parameters(),lr=2e-3,weight_decay=1e-4)
 95    x,y=train; xt,yt=test; x=x.to(device); y=y.to(device); xt=xt.to(device); yt=yt.to(device)
 96    losses=[]; accs=[]; regs=[]
 97    for step in range(160):
 98        idx=torch.randperm(len(x),device=device)[:64]
 99        logits,reg=model(x[idx]); loss=F.cross_entropy(logits,y[idx])+lam*reg
100        opt.zero_grad(); loss.backward(); opt.step()
101        if step in [39,79,159]:
102            with torch.no_grad():
103                z,r=model(xt); acc=(z.argmax(1)==yt).float().mean().item()
104            losses.append(float(loss.item())); accs.append(acc); regs.append(float(reg.item()))
105    return {'val_acc':accs,'train_loss':losses,'reg':regs}
106
107def main():
108    seed_all(); device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
109    try:
110        train=dataset(512,11); test=dataset(256,12)
111        # Keep data on CPU until training; CUDA failures are handled below.
112        out={'device':str(device),'math':math_check()}
113        try:
114            base=run_train(False,0,.5,train,test,device)
115            idea=run_train(True,.01,.5,train,test,device)
116        except Exception as e:
117            if device.type=='cuda':
118                device=torch.device('cpu'); base=run_train(False,0,.5,train,test,device); idea=run_train(True,.01,.5,train,test,device)
119                out['cuda_error']=repr(e)
120            else: raise
121        out.update({'baseline':base,'idea':idea})
122        Path('results.json').write_text(json.dumps(out,indent=2))
123        print(json.dumps(out,indent=2))
124    except RuntimeError as e:
125        if 'out of memory' in str(e).lower() and torch.cuda.is_available():
126            torch.cuda.empty_cache(); print('CUDA OOM; rerun with CPU')
127            raise
128        raise
129if __name__=='__main__': main()