| """Quantize a TinyFaceNet checkpoint and measure accuracy drop. |
| Variants: fp32 | int8_ort (ONNX Runtime static, QOperator, per-channel -- the real deployable path) |
| | int8_w (weight-only fake-quant, per-channel symmetric) | int4_w (weight-only fake-quant, per-channel) | int4_g32 (per-group of 32) |
| Metrics (date-held-out test set): family acc (centroid argmax), open-set acc with THR, mean cosine agreement with fp32 embeddings over the FULL dataset, |
| and argmax agreement with fp32 over the full dataset (statistically more stable than the 59-face test set). |
| usage: python quantize_eval.py <ckpt> <tag> -> analysis/quant_<tag>.json, dataset/<tag>_fp32.onnx, dataset/<tag>_int8.onnx""" |
| import sys, os, json, time, copy, numpy as np, torch, torch.nn.functional as F, onnx, onnxruntime as ort |
| from onnxruntime.quantization import quantize_static, CalibrationDataReader, QuantType, QuantFormat |
| from model import TinyFaceNet |
| torch.set_num_threads(4); CK,TAG=sys.argv[1],sys.argv[2]; THR=0.55 |
| ck=torch.load(CK,weights_only=False); names=list(map(str,ck["names"])); C=ck["centroids"].numpy(); oi=names.index("other") |
| d=np.load("dataset/faces_v2.npz"); X=d["imgs"]; L=np.array(list(map(str,d["labels"]))); DT=d["date"]; S=d["src"] |
| Y=np.array([names.index(l) for l in L]); test=np.array([(x[:6] in ("202511","202605","202606")) for x in DT]) |
| full=np.linspace(0,len(X)-1,1500).astype(int) |
| def prep(x): return ((x.transpose(0,3,1,2).astype(np.float32)/255.)-0.5)/0.5 |
| def base(): m=TinyFaceNet(n_cls=len(names)); m.load_state_dict(ck["model"]); m.eval(); return m |
| def run_torch(m,ix): |
| out=[] |
| with torch.no_grad(): |
| for i in range(0,len(ix),256): out.append(F.normalize(m(torch.from_numpy(prep(X[ix[i:i+256]]))),dim=1).numpy()) |
| return np.concatenate(out) |
| def run_ort(path,ix): |
| s=ort.InferenceSession(path,providers=["CPUExecutionProvider"]); n=s.get_inputs()[0].name; out=[] |
| for i in range(0,len(ix),64): e=s.run(None,{n:prep(X[ix[i:i+64]])})[0]; out.append(e/np.linalg.norm(e,axis=1,keepdims=True)) |
| return np.concatenate(out) |
| def metrics(E,Eref,ix): |
| sim=E[np.isin(ix,np.nonzero(test)[0])]@C.T if False else None |
| te=np.nonzero(test)[0]; Et=E[:len(te)]; yt=Y[te]; sim=Et@C.T; pred=sim.argmax(1); best=sim.max(1) |
| fam=yt!=oi; acc=float((pred[fam]==yt[fam]).mean()) |
| pred_os=np.where(best>=THR,pred,oi); acc_os=float((pred_os==yt).mean()) |
| per={names[c]:round(float((pred[yt==c]==c).mean()),3) for c in range(len(names)) if (yt==c).any()} |
| Ef=E[len(te):]; Rf=Eref[len(te):]; cos=float((Ef*Rf).sum(1).mean()); agree=float(((Ef@C.T).argmax(1)==(Rf@C.T).argmax(1)).mean()) |
| return dict(fam_acc=acc,openset_acc=acc_os,per_class=per,cos_vs_fp32=cos,argmax_agree_vs_fp32=agree) |
| def fake_quant(m,bits,group=None): |
| m=copy.deepcopy(m); qmax=2**(bits-1)-1 |
| for name,p in m.named_parameters(): |
| if p.dim()<2 or name.startswith(("distill","cls")): continue |
| w=p.data; o=w.shape[0]; wf=w.reshape(o,-1) |
| if group and wf.shape[1]%group==0: wg=wf.reshape(o,-1,group); sc=wg.abs().amax(2,keepdim=True).clamp_min(1e-8)/qmax; wf=((wg/sc).round().clamp(-qmax,qmax)*sc).reshape(o,-1) |
| else: sc=wf.abs().amax(1,keepdim=True).clamp_min(1e-8)/qmax; wf=(wf/sc).round().clamp(-qmax,qmax)*sc |
| p.data=wf.reshape(w.shape) |
| return m |
| def lat(fn,n=50): |
| x=np.zeros((1,112,112,3),np.uint8); fn(np.array([0])); t=time.perf_counter() |
| for _ in range(n): fn(np.array([0])) |
| return (time.perf_counter()-t)/n*1000 |
| ix=np.concatenate([np.nonzero(test)[0],full]); res={}; m=base() |
| |
| E0=run_torch(m,ix); res["fp32"]=dict(**metrics(E0,E0,ix),size_kb=round(sum(p.numel() for n,p in m.named_parameters() if not n.startswith(("distill","cls")))*4/1024),lat_ms=round(lat(lambda i:run_torch(m,i)),2),note="PyTorch fp32 (deploy params only)") |
| |
| fp=f"dataset/{TAG}_fp32.onnx"; q8=f"dataset/{TAG}_int8.onnx" |
| torch.onnx.export(m,torch.zeros(1,3,112,112),fp,input_names=["x"],output_names=["emb"],dynamic_axes={"x":{0:"b"},"emb":{0:"b"}},opset_version=17,dynamo=False) |
| E=run_ort(fp,ix); res["fp32_onnx"]=dict(**metrics(E,E0,ix),size_kb=round(os.path.getsize(fp)/1024),lat_ms=round(lat(lambda i:run_ort(fp,i)),2),note="ONNX Runtime fp32") |
| |
| class CR(CalibrationDataReader): |
| def __init__(s): tr=np.nonzero(~test)[0]; s.ix=np.linspace(0,len(tr)-1,200).astype(int); s.tr=tr; s.i=0 |
| def get_next(s): |
| if s.i>=len(s.ix): return None |
| b=s.tr[s.ix[s.i:s.i+8]]; s.i+=8; return {"x":prep(X[b])} |
| quantize_static(fp,q8,CR(),quant_format=QuantFormat.QOperator,per_channel=True,activation_type=QuantType.QUInt8,weight_type=QuantType.QInt8) |
| E=run_ort(q8,ix); res["int8_ort"]=dict(**metrics(E,E0,ix),size_kb=round(os.path.getsize(q8)/1024),lat_ms=round(lat(lambda i:run_ort(q8,i)),2),note="ONNX Runtime static int8, per-channel, 200 calib imgs (deployable)") |
| |
| for tag,bits,g,note in [("int8_w",8,None,"simulated: int8 weights per-channel, fp32 activations"), |
| ("int4_w",4,None,"simulated: int4 weights per-channel (coarse)"), |
| ("int4_g32",4,32,"simulated: int4 weights, group=32 scales (llama.cpp-style)")]: |
| mq=fake_quant(m,bits,g); E=run_torch(mq,ix); core=sum(p.numel() for n,p in m.named_parameters() if not n.startswith(("distill","cls"))) |
| kb=core*bits/8/1024*(1+ (2/g if g else 0.002)) |
| res[tag]=dict(**metrics(E,E0,ix),size_kb=round(kb),lat_ms=None,note=note) |
| json.dump(dict(ckpt=CK,tag=TAG,names=names,n_test=int(test.sum()),n_test_family=int((test&(Y!=oi)).sum()),n_agree=len(full),thr=THR,results=res),open(f"analysis/quant_{TAG}.json","w"),indent=1) |
| for k,v in res.items(): print(f"{k:10s} fam_acc={v['fam_acc']:.3f} openset={v['openset_acc']:.3f} cos={v['cos_vs_fp32']:.4f} agree={v['argmax_agree_vs_fp32']:.3f} size={v['size_kb']}KB lat={v['lat_ms']}ms") |
|
|
