83 lines
6.4 KiB
Python
83 lines
6.4 KiB
Python
#!/usr/bin/env python3
|
|
"""Train a native-rate-code readout; all deployed weights use <=4 magnitudes incl reset255."""
|
|
import os
|
|
os.environ['OPENBLAS_NUM_THREADS']='8'
|
|
os.environ['OMP_NUM_THREADS']='8'
|
|
from pathlib import Path
|
|
import numpy as np,json,time,traceback
|
|
ROOT=Path('/home/user/run/mnist-altai-native'); ROOT.mkdir(exist_ok=True,parents=True)
|
|
TEACH=Path('/home/user/run/mnist-fast'); T=32; HS=8; HT=32; G=8
|
|
start=time.monotonic()
|
|
def save(name,obj):
|
|
p=ROOT/name;q=p.with_suffix(p.suffix+'.tmp');q.write_text(json.dumps(obj,separators=(',',':'))+'\n');q.replace(p)
|
|
def prog(n,label,**kw): save('progress.json',dict(current=n,total=100,done=False,failed=False,label=label,**kw))
|
|
def clustered(a,scale):
|
|
out=np.zeros(a.shape,np.int16)
|
|
for c in range(a.shape[1]):
|
|
for g in range(G):
|
|
ids=np.arange(g*32,(g+1)*32); vals=np.abs(a[ids,c])*scale
|
|
hi=max(2,min(254,int(np.floor(vals.max()+.5))));levels=np.array([0,1,max(2,round(hi*.45)),hi])
|
|
out[ids,c]=levels[np.abs(vals[:,None]-levels).argmin(1)]*np.sign(a[ids,c]).astype(int)
|
|
return out
|
|
|
|
def main():
|
|
rng=np.random.default_rng(20261015);x=np.load(TEACH/'normalized.npy',mmap_mode='r');s=np.load(TEACH/'split.npz');train=s['train'];val=s['validation']
|
|
labels=np.loadtxt('/home/user/run/mnist-native-20261015/MNIST.target',dtype=int)
|
|
teacher=np.load(TEACH/'best.npz');v=teacher['w1']*HS;levels=np.array([0,1,2,4]);w1=(levels[np.abs(np.abs(v)[...,None]-levels).argmin(2)]*np.sign(v)).astype(np.int16);b1=np.floor(teacher['b1']*HS+.5).astype(np.int16)
|
|
bound=T*(np.maximum(w1,0).sum(0)+np.maximum(b1,0));assert bound.max()<32000,(bound.max(),'hidden unsafe')
|
|
feat=np.zeros((70000,256),np.int16);hmin=0;hmax=0
|
|
for first in range(0,70000,512):
|
|
im=x[first:first+512].astype(np.int16);n=len(im);r=np.zeros(im.shape,np.int16);p=np.zeros((n,256),np.float64);c=np.zeros((n,256),np.int16)
|
|
for tick in range(T):
|
|
r+=im;sp=r>=255;r[sp]-=255;p+=sp.astype(np.float64)@w1.astype(np.float64)+b1
|
|
hmin=min(hmin,float(p.min()));hmax=max(hmax,float(p.max()));f=p>=HT;c+=f;p[f]-=HT;np.maximum(p,0,out=p)
|
|
feat[first:first+n]=c;prog(5+45*(first+n)/70000,'native rate hidden features')
|
|
np.save(ROOT/'hidden_counts.npy',feat)
|
|
# A fixed hidden layer inherited from train-only teacher; train readout using train ids only.
|
|
a=feat.astype(np.float32)*(4/T);w=teacher['w2'].copy();b=teacher['b2'].copy();m=[np.zeros_like(w),np.zeros_like(b)];vv=[z.copy() for z in m];updates=0;best=-1;hist=[]
|
|
for epoch in range(35):
|
|
order=rng.permutation(train);lr=.003*(.3 if epoch>=20 else 1)*(.3 if epoch>=30 else 1)
|
|
for k in range(0,len(order),256):
|
|
ix=order[k:k+256];aa=a[ix];logit=aa@w+b;logit-=logit.max(1,keepdims=True);pr=np.exp(logit);pr/=pr.sum(1,keepdims=True);pr[np.arange(len(ix)),labels[ix]]-=1;pr/=len(ix);gr=[aa.T@pr+1e-5*w,pr.sum(0)];updates+=1
|
|
for par,gg,mm,vv0 in zip([w,b],gr,m,vv):
|
|
mm*=.9;mm+=.1*gg;vv0*=.999;vv0+=.001*gg*gg;par-=lr*(mm/(1-.9**updates))/(np.sqrt(vv0/(1-.999**updates))+1e-8)
|
|
acc=float(((a[val]@w+b).argmax(1)==labels[val]).mean());hist.append(dict(epoch=epoch+1,validation=acc))
|
|
if acc>best: best=acc;np.savez(ROOT/'readout-best.npz',w=w,b=b)
|
|
prog(50+25*(epoch+1)/35,'native readout training')
|
|
bestp=np.load(ROOT/'readout-best.npz');w,b=bestp['w'],bestp['b'];candidates=[];chosen=None
|
|
for scale in [4,8,16,32]:
|
|
w2=clustered(w,scale);b2=np.floor(b*scale/4+.5).astype(np.int16)
|
|
# output neurons: sign, group of32 inputs, class. Clock bias goes to group0 only.
|
|
charges=np.empty((70000,2,G,10),np.int64)
|
|
for sign in range(2):
|
|
ww=np.maximum(w2 if sign==0 else -w2,0)
|
|
for g in range(G): charges[:,sign,g]=feat[:,g*32:(g+1)*32].astype(np.int64)@ww[g*32:(g+1)*32].astype(np.int64)
|
|
charges[:,sign,0]+=T*np.maximum(b2 if sign==0 else -b2,0)
|
|
static=max(T*(np.maximum(w2,0).reshape(G,32,10).sum(1).max()+np.abs(b2).max()),T*(np.maximum(-w2,0).reshape(G,32,10).sum(1).max()+np.abs(b2).max()))
|
|
if static>=32000: continue
|
|
for threshold in [4,8,16,32]:
|
|
counts=charges//threshold;scores=counts[:,0].sum(1)-counts[:,1].sum(1);acc=float((scores[val].argmax(1)==labels[val]).mean());entry=dict(scale=scale,threshold=threshold,validation=acc,static_potential_bound=int(static),validation_max_output_count=int(counts[val].max()))
|
|
candidates.append(entry)
|
|
# Prefer accuracy; then larger threshold to reduce official emulator drain time.
|
|
if chosen is None or (acc,threshold)>(chosen['validation'],chosen['threshold']):
|
|
chosen=entry.copy();selected=(w2.copy(),b2.copy(),scores.copy(),counts.copy())
|
|
assert chosen is not None
|
|
w2,b2,scores,counts=selected
|
|
save('selection.json',dict(chosen=chosen,candidates=candidates,hidden_scale=HS,hidden_threshold=HT,active_steps=T,groups=G,train=55000,validation=5000,seed=20261015))
|
|
# Test is evaluated only after selected parameters have been fixed and saved.
|
|
pred=scores[60000:].argmax(1);yy=labels[60000:];conf=np.zeros((10,10),int);np.add.at(conf,(yy,pred),1)
|
|
metrics=dict(test_correct=int((pred==yy).sum()),test_total=10000,test_accuracy=float((pred==yy).mean()),validation=chosen['validation'],hidden_observed_preclip_range=[hmin,hmax],hidden_static_upper=int(bound.max()),output_static_upper=chosen['static_potential_bound'],configuration=chosen,architecture='784 rate inputs ->256 AltAILIF ->160 opponent tiled AltAILIF outputs; host subtract counts and argmax',official_emulator_verified=False,confusion=conf.tolist(),elapsed_seconds=time.monotonic()-start)
|
|
save('metrics-reference.json',metrics);save('history.json',hist)
|
|
# Little-endian model header and signed integer arrays, consumed by the C++ graph builder.
|
|
header=np.array([784,256,10,G,T,HT,chosen['threshold'],max(int(np.abs(b1).max()),int(np.abs(b2).max()))],'<i4')
|
|
with (ROOT/'model.bin').open('wb') as f:
|
|
for ar in [header,w1.astype('<i2'),b1.astype('<i2'),w2.astype('<i2'),b2.astype('<i2')]:f.write(ar.tobytes())
|
|
np.savez(ROOT/'reference.npz',predictions=pred,labels=yy,output_counts=counts[60000:],hidden_counts=feat[60000:],scores=scores[60000:])
|
|
np.savez(ROOT/'model.npz',w1=w1,b1=b1,w2=w2,b2=b2)
|
|
# Normalized pixels, not CPU weighted features, are the sole data given to emulator.
|
|
x.astype(np.uint8).tofile(ROOT/'images.bin');np.array(labels,dtype='<i4').tofile(ROOT/'labels.bin')
|
|
save('progress.json',dict(current=100,total=100,done=True,failed=False,label=f'reference {metrics["test_correct"]}/10000; pending GoldenModel'))
|
|
print(json.dumps(metrics),flush=True)
|
|
if __name__=='__main__':
|
|
try:main()
|
|
except Exception as e:save('progress.json',dict(current=0,total=100,done=False,failed=True,error=str(e)));raise
|