our_knp/pets-knp/train_head_legacy.py

39 lines
3.7 KiB
Python

"""Own offline MLP; export integer AltAILIF hidden rates and native two-neuron readout."""
from pathlib import Path
import numpy as np,json,hashlib,time
R=Path.home()/'run/pets-knp';M=R/'head';M.mkdir(exist_ok=False);rng=np.random.default_rng(20261015)
a=np.load(R/'training-counts.npz');yt=np.tile(np.load(R/'data/train.npz')['labels'],2).astype(np.float32);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)
xi=encode(a['train']);vi=encode(a['validation']);mu=xi.mean(0).astype(np.float32);sd=(xi.std(0)+2).astype(np.float32);x=((xi-mu)/sd).astype(np.float32);v=((vi-mu)/sd).astype(np.float32)
h=256;params=[rng.normal(0,np.sqrt(2/800),(800,h)).astype(np.float32),np.zeros(h,np.float32),rng.normal(0,.05,h).astype(np.float32),np.array(0,np.float32)];mm=[np.zeros_like(p) for p in params];vv=[np.zeros_like(p) for p in params];step=0;history=[];best=None
for epoch in range(1,41):
lr=.0015*(.98**(epoch-1));order=rng.permutation(len(x))
for start in range(0,len(x),256):
ix=order[start:start+256];xx=x[ix];yy=yt[ix];w,b,u,c=params;z=xx@w+b;act=np.clip(z,0,16);mask=(rng.random(act.shape)>.15).astype(np.float32)/.85;hidden=act*mask;scores=hidden@u+c;prob=1/(1+np.exp(-np.clip(scores,-30,30)));dz=(prob-yy)/len(xx)
dh=dz[:,None]*u[None,:]*mask*((z>0)&(z<16));grads=[xx.T@dh+.001*w,dh.sum(0),hidden.T@dz+.001*u,np.array(dz.sum(),np.float32)];step+=1
for k,g in enumerate(grads): mm[k]=.9*mm[k]+.1*g
# Array/scalar Adam, avoid assigning a non-sliceable scalar.
for k,(p,g) in enumerate(zip(params,grads)):
vv[k]=.999*vv[k]+.001*g*g
p-=lr*(mm[k]/(1-.9**step))/(np.sqrt(vv[k]/(1-.999**step))+1e-8)
w,b,u,c=params;pred=(np.clip(v@w+b,0,16)@u+c>0);floatacc=float((pred==yv).mean())
# Fold training standardization into integer input weights; per-neuron scale.
wf=w/sd[:,None];bf=b-(mu/sd)@w;scale=127/np.maximum(abs(wf).max(0),1e-12);wq=np.rint(wf*scale).astype(np.int32);bq=np.rint(bf*scale).astype(np.int32);thr=np.maximum(np.rint(scale/4),1).astype(np.int32)
hidden=np.clip(np.floor(np.maximum(vi.astype(np.int64)@wq+bq,0)/thr),0,64).astype(np.int32)
# Readout refit on train hidden spike counts; offline only. Tune regularisation on validation.
ht=np.clip(np.floor(np.maximum(xi.astype(np.int64)@wq+bq,0)/thr),0,64).astype(np.int32);hm=ht.mean(0);hs=ht.std(0)+.1;zz=(ht-hm)/hs;A=zz.T@zz;rhs=zz.T@(yt*2-1)
epochbest=None
for lam in (10,100,1000):
sol=np.linalg.solve(A+lam*np.eye(h),rhs);ww=sol/hs;bb=(yt*2-1).mean()-hm@ww;ss=127/abs(ww).max();uw=np.rint(ww*ss).astype(np.int32);ub=int(round(bb*ss));margin=hidden.astype(np.int64)@uw+ub;acc=float(((margin>0)==yv).mean())
if epochbest is None or acc>epochbest[0]:epochbest=(acc,uw,ub,lam)
row={'epoch':epoch,'float_validation':floatacc,'native_quantized_validation':epochbest[0],'lambda':epochbest[3]};history.append(row);print(row,flush=True)
if best is None or epochbest[0]>best:
best=epochbest[0];np.savez(M/'model.npz',w=wq,b=bq,thresholds=thr,readout_w=epochbest[1],readout_b=np.array([epochbest[2]],np.int32));(M/'selection.json').write_text(json.dumps(dict(row,encoding='root_l2',hidden=h,max_hidden_count=64),indent=2))
(M/'progress.json').write_text(json.dumps({'current':epoch,'total':40,'done':epoch==40}))
(M/'history.json').write_text(json.dumps(history,indent=2))
# Source artifact export; no test read by this script.
m=np.load(M/'model.npz');
for key,fn in [('w','hidden_w.bin'),('b','hidden_b.bin'),('thresholds','hidden_thresholds.bin'),('readout_w','weights.bin'),('readout_b','bias.bin')]:m[key].astype(np.int32).tofile(M/fn)
vi.tofile(R/'validation-encoded.bin');print('BEST',best,flush=True)