Spiking readout (off-chip ridge, AltAI-like integer inference) and robustness check
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
int
parent
7add19dcae
commit
b8ea9dc6ce
2 changed files with 165 additions and 0 deletions
51
readout/robustness.py
Normal file
51
readout/robustness.py
Normal file
|
|
@ -0,0 +1,51 @@
|
||||||
|
"""Robustness of the spiking readout: distort raw validation pixels, compute stroke features, run the readout.
|
||||||
|
|
||||||
|
usage: python robustness.py <raw.bin> <raw.target> <feat_train.bin> <feat_train.target> --train N [--preset coarse --scale 0.3842]
|
||||||
|
"""
|
||||||
|
import argparse
|
||||||
|
import sys
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
sys.path.insert(0, "/mnt/c/knp/features")
|
||||||
|
sys.path.insert(0, "/mnt/c/knp/augment")
|
||||||
|
import stroke_features as sf # noqa: E402
|
||||||
|
import spiking_readout as sr # noqa: E402
|
||||||
|
from tta_sim import grow, shift, thin # noqa: E402
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
ap = argparse.ArgumentParser()
|
||||||
|
ap.add_argument("raw")
|
||||||
|
ap.add_argument("raw_labels")
|
||||||
|
ap.add_argument("feat")
|
||||||
|
ap.add_argument("feat_labels")
|
||||||
|
ap.add_argument("--train", type=int, required=True)
|
||||||
|
ap.add_argument("--eval-first", type=int, default=30000)
|
||||||
|
ap.add_argument("--eval", type=int, default=10000)
|
||||||
|
ap.add_argument("--preset", default="coarse")
|
||||||
|
ap.add_argument("--scale", type=float, default=0.3842)
|
||||||
|
ap.add_argument("--thr", type=float, default=16)
|
||||||
|
ap.add_argument("--lam", type=float, default=10.0)
|
||||||
|
ap.add_argument("--save")
|
||||||
|
a = ap.parse_args()
|
||||||
|
sf.use_preset(a.preset)
|
||||||
|
sf.P["scale"] = a.scale
|
||||||
|
width = 784
|
||||||
|
xf = np.fromfile(a.feat, dtype=np.uint8).reshape(-1, width)[:a.train]
|
||||||
|
yf = np.array([int(v) for v in open(a.feat_labels).read().split()])[:a.train]
|
||||||
|
w = sr.train_ridge(xf, yf, 10, a.lam)
|
||||||
|
wq, scale = sr.quantize(w)
|
||||||
|
raw = np.fromfile(a.raw, dtype=np.uint8).reshape(-1, 28, 28)[a.eval_first:a.eval_first + a.eval].astype(np.int32)
|
||||||
|
y = np.array([int(v) for v in open(a.raw_labels).read().split()])[a.eval_first:a.eval_first + a.eval]
|
||||||
|
row = []
|
||||||
|
for name, v in [("as is", raw), ("shift 2px", shift(raw, 0, 2)), ("thicker 1px", grow(raw)), ("thinner 1px", thin(raw))]:
|
||||||
|
p = sr.spiking_predict(sf.features(v.clip(0, 255).astype(np.uint8)), wq, a.thr)
|
||||||
|
row.append(f"{name} {100 * (p == y).mean():5.1f}%")
|
||||||
|
print(" | ".join(row))
|
||||||
|
if a.save:
|
||||||
|
np.savez(a.save, w=wq, thr=a.thr, scale=scale, preset=a.preset, feat_scale=a.scale)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
114
readout/spiking_readout.py
Normal file
114
readout/spiking_readout.py
Normal file
|
|
@ -0,0 +1,114 @@
|
||||||
|
"""Spiking readout: weights trained off-chip (ridge regression), inference by AltAI-like integer neurons.
|
||||||
|
|
||||||
|
Input neurons fire with the same incrementing raster as mnist-learn (value/255 per step, 10 active steps
|
||||||
|
of 15). One output neuron per class integrates integer weights of the spikes it receives, fires at
|
||||||
|
`thr` and resets to 0; the potential is clipped at 0 like the AltAI negative threshold. A bias input
|
||||||
|
fires on every active step. The answer is the class whose neuron fired most (earliest on ties).
|
||||||
|
|
||||||
|
usage: python spiking_readout.py <data.bin> <labels> --width 784 --train 30000 --eval-first 30000 --eval 10000
|
||||||
|
[--save weights.npz] [--thr 50,100,200]
|
||||||
|
"""
|
||||||
|
import argparse
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
STEPS, ACTIVE = 15, 10
|
||||||
|
|
||||||
|
|
||||||
|
def spike_counts(x):
|
||||||
|
"""x: [n, width] uint8 → spikes per input over ACTIVE steps (raster from rest)."""
|
||||||
|
st = np.zeros(x.shape, dtype=np.float32)
|
||||||
|
inc = x.astype(np.float32) * np.float32(1 / 255.0)
|
||||||
|
cnt = np.zeros(x.shape, dtype=np.int16)
|
||||||
|
for _ in range(ACTIVE):
|
||||||
|
st += inc
|
||||||
|
f = st >= 1.0
|
||||||
|
cnt += f
|
||||||
|
st[f] -= 1.0
|
||||||
|
return cnt
|
||||||
|
|
||||||
|
|
||||||
|
def raster(x):
|
||||||
|
"""x: [n, width] uint8 → [ACTIVE, n, width] bool."""
|
||||||
|
st = np.zeros(x.shape, dtype=np.float32)
|
||||||
|
inc = x.astype(np.float32) * np.float32(1 / 255.0)
|
||||||
|
out = np.zeros((ACTIVE,) + x.shape, dtype=bool)
|
||||||
|
for t in range(ACTIVE):
|
||||||
|
st += inc
|
||||||
|
out[t] = st >= 1.0
|
||||||
|
st[out[t]] -= 1.0
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
def train_ridge(x, y, classes, lam):
|
||||||
|
f = np.hstack([spike_counts(x).astype(np.float32), np.full((len(x), 1), ACTIVE, np.float32)])
|
||||||
|
t = np.eye(classes, dtype=np.float32)[y] * 2 - 1
|
||||||
|
return np.linalg.solve(f.T @ f + lam * np.eye(f.shape[1], dtype=np.float32), f.T @ t)
|
||||||
|
|
||||||
|
|
||||||
|
def quantize(w, wmax=255):
|
||||||
|
scale = wmax / np.abs(w).max()
|
||||||
|
return np.round(w * scale).astype(np.int32), scale
|
||||||
|
|
||||||
|
|
||||||
|
def spiking_predict(x, wq, thr, batch=2000):
|
||||||
|
preds = []
|
||||||
|
for i in range(0, len(x), batch):
|
||||||
|
r = raster(x[i:i + batch])
|
||||||
|
n = r.shape[1]
|
||||||
|
pot = np.zeros((n, wq.shape[1]), dtype=np.int64)
|
||||||
|
count = np.zeros_like(pot)
|
||||||
|
first = np.full(pot.shape, STEPS, dtype=np.int64)
|
||||||
|
bias = wq[-1]
|
||||||
|
for t in range(STEPS):
|
||||||
|
if t < ACTIVE:
|
||||||
|
pot += r[t].astype(np.int64) @ wq[:-1] + bias
|
||||||
|
fired = pot >= thr
|
||||||
|
count += fired
|
||||||
|
first = np.where(fired & (first == STEPS), t, first)
|
||||||
|
pot[fired] = 0
|
||||||
|
pot = np.maximum(pot, 0)
|
||||||
|
# most spikes; ties broken by earliest first spike; no spikes → -1
|
||||||
|
score = count * (STEPS + 1) - first
|
||||||
|
p = score.argmax(1)
|
||||||
|
p[count.max(1) == 0] = -1
|
||||||
|
preds.append(p)
|
||||||
|
return np.concatenate(preds)
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
ap = argparse.ArgumentParser()
|
||||||
|
ap.add_argument("data")
|
||||||
|
ap.add_argument("labels")
|
||||||
|
ap.add_argument("--width", type=int, default=784)
|
||||||
|
ap.add_argument("--classes", type=int, default=10)
|
||||||
|
ap.add_argument("--train", type=int, required=True)
|
||||||
|
ap.add_argument("--eval-first", type=int, required=True)
|
||||||
|
ap.add_argument("--eval", type=int, required=True)
|
||||||
|
ap.add_argument("--lam", type=float, default=10.0)
|
||||||
|
ap.add_argument("--thr", default="")
|
||||||
|
ap.add_argument("--save")
|
||||||
|
a = ap.parse_args()
|
||||||
|
x = np.fromfile(a.data, dtype=np.uint8).reshape(-1, a.width)
|
||||||
|
y = np.array([int(v) for v in open(a.labels).read().split()])
|
||||||
|
w = train_ridge(x[:a.train], y[:a.train], a.classes, a.lam)
|
||||||
|
xe, ye = x[a.eval_first:a.eval_first + a.eval], y[a.eval_first:a.eval_first + a.eval]
|
||||||
|
fe = np.hstack([spike_counts(xe).astype(np.float32), np.full((len(xe), 1), ACTIVE, np.float32)])
|
||||||
|
print(f"float linear on spike counts: {100 * ((fe @ w).argmax(1) == ye).mean():.2f}%")
|
||||||
|
wq, scale = quantize(w)
|
||||||
|
ref = np.abs(wq).max()
|
||||||
|
thrs = [float(t) for t in a.thr.split(",")] if a.thr else [ref * k for k in (0.5, 1, 2, 4, 8)]
|
||||||
|
best = None
|
||||||
|
for thr in thrs:
|
||||||
|
p = spiking_predict(xe, wq, thr)
|
||||||
|
acc = 100 * (p == ye).mean()
|
||||||
|
print(f"spiking, thr {thr:g}: correct {acc:.2f}% silent {100 * (p < 0).mean():.2f}%")
|
||||||
|
if best is None or acc > best[0]:
|
||||||
|
best = (acc, thr)
|
||||||
|
if a.save:
|
||||||
|
np.savez(a.save, w=wq, thr=best[1], scale=scale, width=a.width)
|
||||||
|
print("saved", a.save, "thr", best[1])
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
Loading…
Add table
Reference in a new issue