| import os, numpy as np, torch, torch.nn as nn, torch.nn.functional as F, time, json, collections, sys |
| from model import TinyFaceNet |
| torch.manual_seed(0); np.random.seed(0); torch.set_num_threads(int(os.environ.get("THREADS","10"))) |
| import os; d=np.load(os.environ.get("DATA","dataset/faces.npz")); X=d["imgs"]; T=d["embs"]; L=d["labels"]; S=d["src"]; DT=d["date"] |
| names=sorted(set(L)); idx={n:i for i,n in enumerate(names)}; Y=np.array([idx[l] for l in L]) |
| |
| test=np.array([(x[:6] in ("202511","202605","202606")) for x in DT]); train=~test |
| print("classes",names); print("train",train.sum(),"test",test.sum()) |
| print("train per class",dict(collections.Counter(L[train]))); print("test per class",dict(collections.Counter(L[test]))) |
| EPOCHS=int(sys.argv[1]) if len(sys.argv)>1 else 40; BS=128 |
| def aug(x): |
| x=torch.from_numpy(x).permute(0,3,1,2).float()/255. |
| B=x.shape[0] |
| flip=torch.rand(B)<0.5; x[flip]=x[flip].flip(-1) |
| |
| th=torch.zeros(B,2,3); sc=1+0.15*(torch.rand(B)*2-1); th[:,0,0]=sc; th[:,1,1]=sc; th[:,:,2]=0.12*(torch.rand(B,2)*2-1) |
| x=F.grid_sample(x,F.affine_grid(th,x.shape,align_corners=False),padding_mode="border",align_corners=False) |
| |
| x=x*(1+0.4*(torch.rand(B,1,1,1)-0.5))+0.2*(torch.rand(B,1,1,1)-0.5) |
| |
| g=torch.rand(B)<0.4 |
| if g.any(): x[g]=x[g].clamp(1e-3,1)**(1/(0.3+0.7*torch.rand(int(g.sum()),1,1,1)))*(0.5+0.5*torch.rand(int(g.sum()),1,1,1)) |
| cc=torch.rand(B)<0.3 |
| if cc.any(): mu=x[cc].mean((1,2,3),keepdim=True); x[cc]=(x[cc]-mu)*(0.4+0.6*torch.rand(int(cc.sum()),1,1,1))+mu |
| nz=torch.rand(B)<0.4 |
| if nz.any(): x[nz]=x[nz]+torch.randn_like(x[nz])*(0.08*torch.rand(int(nz.sum()),1,1,1)) |
| gray=torch.rand(B)<0.15; x[gray]=x[gray].mean(1,keepdim=True).expand(-1,3,-1,-1) |
| blur=torch.rand(B)<0.2 |
| if blur.any(): x[blur]=F.avg_pool2d(x[blur],3,1,1) |
| |
| for i in torch.nonzero(torch.rand(B)<0.3).flatten(): |
| cy,cx=np.random.randint(0,112,2); x[i,:,max(0,cy-16):cy+16,max(0,cx-16):cx+16]=0 |
| return (x.clamp(0,1)-0.5)/0.5 |
| def prep(x): return (torch.from_numpy(x).permute(0,3,1,2).float()/255.-0.5)/0.5 |
| m=TinyFaceNet(n_cls=len(names)); opt=torch.optim.AdamW(m.parameters(),2e-3,weight_decay=5e-4) |
| tr=np.nonzero(train)[0]; te=np.nonzero(test)[0]; Tt=torch.from_numpy(T) |
| |
| OTHER_CAP=600 |
| oth=tr[L[tr]=="other"]; fam=tr[L[tr]!="other"] |
| oj=oth[np.char.startswith(S[oth],"jpg:")]; ov=oth[~np.char.startswith(S[oth],"jpg:")] |
| need=max(0,OTHER_CAP-len(oj)); ov=ov[np.linspace(0,len(ov)-1,min(need,len(ov))).astype(int)] if need else ov[:0] |
| tr=np.concatenate([fam,oj[:OTHER_CAP],ov]) |
| print("train after other-cap",len(tr),dict(collections.Counter(L[tr]))) |
| |
| cnt=collections.Counter(Y[tr]); w=np.array([1/np.sqrt(cnt[y]) for y in Y[tr]]); w/=w.sum() |
| steps=EPOCHS*(len(tr)//BS+1); sched=torch.optim.lr_scheduler.OneCycleLR(opt,2e-3,total_steps=steps,pct_start=0.1) |
| def embed(ix): |
| m.eval(); out=[] |
| with torch.no_grad(): |
| for i in range(0,len(ix),256): out.append(F.normalize(m(prep(X[ix[i:i+256]])),dim=1)) |
| return torch.cat(out) |
| def evaluate(): |
| Etr=embed(tr); Ete=embed(te); ytr=torch.from_numpy(Y[tr]); yte=torch.from_numpy(Y[te]) |
| C=torch.stack([F.normalize(Etr[ytr==c].mean(0),dim=0) if (ytr==c).any() else torch.zeros(Etr.shape[1]) for c in range(len(names))]) |
| sim=Ete@C.T; pred=sim.argmax(1); acc=(pred==yte).float().mean().item() |
| per={str(names[c]):round((pred[yte==c]==c).float().mean().item(),3) for c in range(len(names)) if (yte==c).any()} |
| fam=yte!=idx["other"]; facc=(pred[fam]==yte[fam]).float().mean().item() if fam.any() else 0. |
| return facc,per,C |
| best=0 |
| for ep in range(EPOCHS): |
| m.train(); ep_ix=np.random.choice(tr,size=len(tr),replace=True,p=w); t0=time.time(); tl=0 |
| for i in range(0,len(ep_ix),BS): |
| b=ep_ix[i:i+BS]; x=aug(X[b]); y=torch.from_numpy(Y[b]); t=Tt[b] |
| e=m(x); lc=F.cross_entropy(m.cls(e),y,label_smoothing=0.1); ld=1-F.cosine_similarity(m.distill(e),t).mean() |
| loss=lc+2.0*ld; opt.zero_grad(); loss.backward(); opt.step(); sched.step(); tl+=loss.item()*len(b) |
| acc,per,C=evaluate() |
| print(f"ep{ep} loss={tl/len(tr):.3f} fam_acc={acc:.4f} {per} {time.time()-t0:.0f}s",flush=True) |
| if acc>=best: best=acc; torch.save({"model":m.state_dict(),"names":names,"centroids":C},os.environ.get("OUT","dataset/best.pt")) |
| print("BEST",best) |
|
|
