our_knp/pets-knp/adapt_head.py

24 lines
2.2 KiB
Python

"""Repair export scales to real uint16 thresholds; refit only from training data."""
from pathlib import Path
import numpy as np,json,hashlib
R=Path.home()/'run/pets-knp';old=R/'head';M=R/'head-v2';M.mkdir(exist_ok=False);a=np.load(old/'model.npz');w=a['w'].copy();b=a['b'].copy();t=a['thresholds'].copy();dead=t<1;large=t>65535
# Negative-overflow thresholds had zero rate in the old evaluator: prune these units.
w[:,dead]=0;b[dead]=0;t[dead]=1
for h in np.flatnonzero(large):
scale=65535/float(t[h]);w[:,h]=np.rint(w[:,h]*scale).astype(np.int32);b[h]=int(round(float(b[h])*scale));t[h]=65535
assert t.min()>=1 and t.max()<=65535
z=np.load(R/'training-counts.npz');yt=np.tile(np.load(R/'data/train.npz')['labels'],2);yv=np.load(R/'data/val.npz')['labels']
def encode(x):
q=np.sqrt(x.astype(np.float64));q*=2048/np.maximum(np.sqrt((q*q).sum(1,keepdims=True)),1);return np.clip(np.rint(q),0,640).astype(np.int32)
def hidden(x):return np.clip(np.floor(np.maximum(encode(x).astype(np.float64)@w.astype(np.float64)+b,0)/t),0,64).astype(np.int32)
ht=hidden(z['train']);hv=hidden(z['validation']);mu=ht.mean(0);sd=ht.std(0)+.1;xx=(ht-mu)/sd;A=xx.T@xx;rhs=xx.T@(yt*2.-1);best=None;rows=[]
for lam in (1,10,100,1000):
sol=np.linalg.solve(A+lam*np.eye(256),rhs);ww=sol/sd;bb=(yt*2.-1).mean()-mu@ww
for limit in (31,63,127):
scale=limit/abs(ww).max();qw=np.rint(ww*scale).astype(np.int32);qb=int(round(bb*scale));m=hv.astype(np.int64)@qw+qb;acc=float(((m>0)==yv).mean());row={'lambda':lam,'weight_limit':limit,'validation_accuracy':acc,'ties':int((m==0).sum())};rows.append(row)
if best is None or acc>best[0]:best=(acc,dict(row),qw,qb)
_,row,u,c=best
np.savez(M/'model.npz',w=w,b=b,thresholds=t,readout_w=u,readout_b=np.array([c],np.int32))
for key,value in [('hidden_w',w),('hidden_b',b),('hidden_thresholds',t),('weights',u),('bias',np.array([c],np.int32))]:value.astype(np.int32).tofile(M/(key+'.bin'))
row.update(encoding='root_l2',hidden=256,max_hidden_count=64,candidates=rows,source_epoch=json.loads((old/'selection.json').read_text())['epoch'],pruned_export_overflow_units=int(dead.sum()),rescaled_large_threshold_units=int(large.sum()),test_used=False)
(M/'selection.json').write_text(json.dumps(row,indent=2));print(json.dumps(row),flush=True)