1import os, sys, json, random, time
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, make_model, train_model, make_report, sweep_baseline, evaluate
  8
  9TRACK='vision'; MODEL='cnn_small'; EPOCHS=3; NTR=400; NTE=200; BATCH=128
 10LR_GRID=[1e-3, 3e-3, 1e-2]
 11SEEDS=tuple(range(8))
 12
 13
 14def seed_all(s):
 15    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 16    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
 17
 18
 19def interval_conv(conv, lo, hi):
 20    w=conv.weight.detach(); b=conv.bias.detach() if conv.bias is not None else 0
 21    wp=torch.clamp(w,min=0); wn=torch.clamp(w,max=0)
 22    return (nn.functional.conv2d(lo,wp,bias=None,padding=conv.padding)+
 23            nn.functional.conv2d(hi,wn,bias=None,padding=conv.padding)+b.view(1,-1,1,1),
 24            nn.functional.conv2d(hi,wp,bias=None,padding=conv.padding)+
 25            nn.functional.conv2d(lo,wn,bias=None,padding=conv.padding)+b.view(1,-1,1,1))
 26
 27
 28def prune_signature(net, x, eps=.02, tau=0.0):
 29    """Sound interval sign classification for the two CNN ReLU blocks.
 30    It evaluates the trained model on x and an independent perturbation sample.
 31    """
 32    m=net.net; lo=torch.clamp(x-eps,0,1); hi=torch.clamp(x+eps,0,1)
 33    fixed=unstable=viol=total=0; max_err=0.0; zlo,zhi=lo,hi
 34    for layer in m:
 35        if isinstance(layer,nn.Conv2d):
 36            L,U=interval_conv(layer,zlo,zhi)
 37            # retain spatial neuron instances; exact fixed-sign substitution
 38            active=L>=0; inactive=U<=0; unst=~(active|inactive)
 39            fixed += int(active.sum()+inactive.sum()); unstable += int(unst.sum()); total += int(L.numel())
 40            # independent point check, using the same certified bounds
 41            q=lo + (hi-lo)*torch.rand_like(lo)
 42            a=layer(q if zlo is lo else zlo)
 43            exact=torch.relu(a)
 44            substituted=torch.where(active,a,torch.where(inactive,torch.zeros_like(a),torch.relu(a)))
 45            max_err=max(max_err,float((exact-substituted).abs().max()))
 46            # propagate interval through the ReLU and pooling
 47            zlo,zhi=torch.relu(L),torch.relu(U)
 48        elif isinstance(layer,nn.ReLU):
 49            pass
 50        elif isinstance(layer,nn.MaxPool2d):
 51            zlo=layer(zlo); zhi=layer(zhi)
 52        elif isinstance(layer,nn.Flatten):
 53            zlo=layer(zlo); zhi=layer(zhi)
 54        elif isinstance(layer,nn.Linear):
 55            break
 56    # Behavioural test: exact substitutions must match ordinary predictions
 57    with torch.no_grad():
 58        normal=net(x).argmax(1)
 59    return {'fixed_sign_fraction': fixed/max(total,1), 'fixed_sign':fixed,
 60            'unstable':unstable, 'total_relu_instances':total,
 61            'max_fixed_substitution_error':max_err,
 62            'prediction_count':int(len(normal)), 'tau':tau,
 63            'predicted_reduction_fraction':fixed/max(total,1)}
 64
 65
 66def train_one(cfg, seed, collect=False):
 67    seed_all(seed); d=get_dataset(TRACK,seed,n_train=NTR,n_test=NTE)
 68    net=make_model(MODEL,d['input_shape'],d['out_dim'])
 69    net, metric, hist=train_model(net,d,epochs=EPOCHS,lr=float(cfg['lr']),batch=BATCH,log=lambda *_:None)
 70    if net is None: return float('inf')
 71    if collect:
 72        net = net.cpu()
 73        sig=prune_signature(net,d['xte'][:min(32,len(d['xte']))])
 74        return metric,sig
 75    return metric
 76
 77
 78def main():
 79    # Baseline sweep over exactly the union of all idea learning rates.
 80    base=sweep_baseline(lambda cfg: (lambda seed: train_one(cfg,seed)),
 81                        [{'lr':lr} for lr in LR_GRID], seeds=(0,1,2,3))
 82    # Re-evaluate the selected baseline and all three idea settings on paired seeds.
 83    base_full=evaluate(lambda s: train_one(base['best_cfg'],s), seeds=SEEDS)
 84    base['full']=base_full
 85    idea_results=[]; signatures=[]
 86    for cfg in [{'lr':lr} for lr in LR_GRID]:
 87        r=evaluate(lambda s,cfg=cfg: train_one(cfg,s), seeds=SEEDS)
 88        idea_results.append({'cfg':cfg,'result':r})
 89    best=min(idea_results,key=lambda z:z['result']['mean'])
 90    for s in SEEDS:
 91        _,sig=train_one(best['cfg'],s,collect=True); signatures.append(sig)
 92    sig={k:float(np.mean([x[k] for x in signatures])) for k in ('fixed_sign_fraction','max_fixed_substitution_error','predicted_reduction_fraction')}
 93    sig.update({'fixed_sign_mean':float(np.mean([x['fixed_sign'] for x in signatures])),
 94                'unstable_mean':float(np.mean([x['unstable'] for x in signatures])),
 95                'total_relu_instances_mean':float(np.mean([x['total_relu_instances'] for x in signatures])),
 96                'confirmed': bool(sig['max_fixed_substitution_error'] <= 1e-7 and sig['predicted_reduction_fraction'] > 0)})
 97    rep=make_report(TRACK,MODEL,base,best['result'],extra={'mechanism_signature':sig,
 98        'idea_sweep':idea_results,
 99        'protocol_note':'Verifier pruning is inference-time only; training and standard test metric are unchanged systems.'})
100    rep['experiment_settings']={'epochs':EPOCHS,'n_train':NTR,'n_test':NTE,'batch':BATCH,'lr_union':LR_GRID,'seeds':list(SEEDS)}
101    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
102    print(json.dumps(rep,indent=2))
103
104if __name__=='__main__': main()