| |
| """Reproducible GPU/CPU scaling benchmark for Hilbert Operator Intelligence.""" |
| from __future__ import annotations |
| import argparse, json, os, random, time |
| from pathlib import Path |
| import numpy as np |
| import torch |
| from torch import nn |
| from torch.utils.data import DataLoader, Subset, TensorDataset |
| ROOT=Path(__file__).resolve().parents[1] |
| import sys; sys.path.insert(0,str(ROOT)) |
| from hilbert_mi.scalable import build_model |
|
|
| def seed_all(s): random.seed(s); np.random.seed(s); torch.manual_seed(s); torch.cuda.manual_seed_all(s) |
| def synthetic(n, classes, size, seed): |
| g=torch.Generator().manual_seed(seed); y=torch.randint(classes,(n,),generator=g); x=torch.randn(n,3,size,size,generator=g)*.15 |
| |
| for i in range(n): |
| k=int(y[i]); x[i,:, (k*3)%size::10, :] += .6; x[i,:, :, (k*5)%size::11] += .35 |
| return TensorDataset(x,y) |
| def datasets(name, root, ntrain, ntest, size, seed): |
| if name=="synthetic": return synthetic(ntrain,10,size,seed), synthetic(ntest,10,size,seed+1),10 |
| try: |
| import torchvision.transforms as T; import torchvision.datasets as D |
| except ImportError as e: raise SystemExit("Install benchmark extras: pip install -e '.[benchmark]'") from e |
| tf=T.Compose([T.Resize((size,size)),T.ToTensor()]) |
| cls={"cifar10":D.CIFAR10,"cifar100":D.CIFAR100}[name]; nc=10 if name=="cifar10" else 100 |
| tr=cls(root=root,train=True,download=True,transform=tf); te=cls(root=root,train=False,download=True,transform=tf) |
| if ntrain and ntrain<len(tr): tr=Subset(tr,range(ntrain)) |
| if ntest and ntest<len(te): te=Subset(te,range(ntest)) |
| return tr,te,nc |
| def evaluate(m,loader,dev): |
| m.eval(); ok=n=0; loss=0.; ce=nn.CrossEntropyLoss(reduction="sum") |
| with torch.no_grad(): |
| for x,y in loader: |
| x,y=x.to(dev),y.to(dev); z=m(x); loss+=ce(z,y).item(); ok+=(z.argmax(1)==y).sum().item(); n+=len(y) |
| return {"accuracy":ok/n,"nll":loss/n} |
| def main(): |
| p=argparse.ArgumentParser(); p.add_argument("--config",required=True); p.add_argument("--dataset",default="cifar10",choices=["cifar10","cifar100","synthetic"]); p.add_argument("--epochs",type=int,default=None); p.add_argument("--batch-size",type=int,default=None); p.add_argument("--train-samples",type=int,default=0); p.add_argument("--test-samples",type=int,default=0); p.add_argument("--image-size",type=int,default=32); p.add_argument("--seed",type=int,default=7); p.add_argument("--output",default="results/run.json"); p.add_argument("--amp",action="store_true"); a=p.parse_args() |
| cfg=json.load(open(a.config)); seed_all(a.seed); dev=torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| tr,te,nc=datasets(a.dataset,str(ROOT/"data"),a.train_samples,a.test_samples,a.image_size,a.seed) |
| bs=a.batch_size or cfg.get("batch_size",128); epochs=a.epochs or cfg.get("epochs",20) |
| dl=DataLoader(tr,batch_size=bs,shuffle=True,num_workers=cfg.get("workers",2),pin_memory=dev.type=="cuda"); tl=DataLoader(te,batch_size=bs,shuffle=False,num_workers=cfg.get("workers",2),pin_memory=dev.type=="cuda") |
| model=build_model(cfg,nc).to(dev); opt=torch.optim.AdamW(model.parameters(),lr=cfg.get("lr",3e-4),weight_decay=cfg.get("weight_decay",1e-4)); ce=nn.CrossEntropyLoss(); scaler=torch.amp.GradScaler("cuda",enabled=a.amp and dev.type=="cuda") |
| t0=time.perf_counter(); seen=0 |
| for ep in range(epochs): |
| model.train() |
| for x,y in dl: |
| x,y=x.to(dev,non_blocking=True),y.to(dev,non_blocking=True); opt.zero_grad(set_to_none=True) |
| with torch.autocast(device_type=dev.type,enabled=a.amp and dev.type=="cuda"): |
| loss=ce(model(x),y) |
| scaler.scale(loss).backward(); scaler.step(opt); scaler.update(); seen+=len(y) |
| if dev.type=="cuda": torch.cuda.synchronize() |
| sec=time.perf_counter()-t0; met=evaluate(model,tl,dev) |
| result={"benchmark":"hoi-image-v1","dataset":a.dataset,"config":Path(a.config).stem,"seed":a.seed,"epochs":epochs,"train_samples":len(tr),"test_samples":len(te),"image_size":a.image_size,"params":model.param_count,"device":str(dev),"gpu":torch.cuda.get_device_name(0) if dev.type=="cuda" else None,"torch":torch.__version__,"amp":bool(a.amp and dev.type=="cuda"),"train_seconds":sec,"samples_per_second":seen/sec,"peak_gpu_memory_gb":torch.cuda.max_memory_allocated()/1e9 if dev.type=="cuda" else 0.0,**met} |
| out=ROOT/a.output; out.parent.mkdir(parents=True,exist_ok=True); out.write_text(json.dumps(result,indent=2)+"\n"); print(json.dumps(result,indent=2)) |
| if __name__=="__main__": main() |
|
|
