152 lines
7.1 KiB
Python
152 lines
7.1 KiB
Python
"""Exact step-by-step spiking simulation of the object pipeline: pixels → layer 1 → readout neuron.
|
||
|
||
Pixels (3×32×32, value/255 per step, 10 active steps of 15) drive layer-1 neurons: one neuron per learned
|
||
prototype k and patch position (27×27), integer weights over its 6×6×3 patch. Patch centring, ZCA and the
|
||
prototype fold into one weight vector; the prototype offset and alpha become a per-step bias (the bias of a
|
||
neuron that also receives a constant input). A neuron adds its input each step, fires at `thr1`, resets to 0
|
||
and never goes below 0 (AltAI). Its spikes go, with delay 1, to the readout neuron "dog − cat" with a weight
|
||
shared by prototype k and image quadrant (this replaces pooling neurons, which could fire only once a step).
|
||
The readout weights are trained off-chip (ridge) on layer-1 spike counts of the training images.
|
||
|
||
usage: ~/anvenv/bin/python spiking_pipeline.py --model /home/int/run/cats-dogs/features-v3/kmeans_model_k800.npz
|
||
"""
|
||
import argparse
|
||
import json
|
||
import sys
|
||
import time
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
|
||
sys.path.insert(0, "/mnt/c/knp/objects")
|
||
sys.path.insert(0, "/mnt/c/knp/readout")
|
||
import kmeans_features as kf # noqa: E402
|
||
|
||
STEPS, ACTIVE, PATCH, SIDE = 15, 10, 6, 27
|
||
WSCALE = 127 # layer-1 weights are integers in ±WSCALE
|
||
|
||
|
||
def fold(model, alpha):
|
||
"""Float layer-1 weights [108, K] and bias [K] for raw patches in [0, 1]: act = p @ W + b."""
|
||
n = PATCH * PATCH * 3
|
||
centre = (np.eye(n) - 1 / n) * 4 # kf.normalise with CONTRAST_DIVIDE=False
|
||
w = centre @ model["zca"] @ model["d"].T
|
||
b = -(model["mean"] @ model["zca"] @ model["d"].T) - alpha
|
||
return w.astype(np.float32), b.astype(np.float32)
|
||
|
||
|
||
def quant_layer1(w, b):
|
||
s = WSCALE / np.abs(w).max()
|
||
# Over ACTIVE steps the pixel spikes sum to ~ACTIVE * p, so the bias is added on each active step.
|
||
return np.rint(w * s).astype(np.int32), np.rint(b * s).astype(np.int32), s
|
||
|
||
|
||
def raster(x):
|
||
"""x uint8 [n, 3, 32, 32] → bool [ACTIVE, n, 3, 32, 32], incrementing converter from rest."""
|
||
st = np.zeros(x.shape, np.float32)
|
||
inc = x.astype(np.float32) * np.float32(1 / 255.0)
|
||
out = np.zeros((ACTIVE,) + x.shape, bool)
|
||
for t in range(ACTIVE):
|
||
st += inc
|
||
out[t] = st >= 1.0
|
||
st[out[t]] -= 1.0
|
||
return out
|
||
|
||
|
||
def layer1_counts(x, wq, bq, thr1, batch=40, per_step=False):
|
||
"""Spike counts of layer-1 neurons pooled per (quadrant, k): [n, 4 * K] (or [n, STEPS, 4 * K] with
|
||
per_step); also total spikes per image."""
|
||
pos = range(SIDE)
|
||
edges = [0, 14, 27]
|
||
res, total = [], []
|
||
for i in range(0, len(x), batch):
|
||
r = raster(x[i:i + batch])
|
||
n = r.shape[1]
|
||
pot = np.zeros((n, SIDE * SIDE, wq.shape[1]), np.int32)
|
||
cnt = np.zeros_like(pot, dtype=np.int16)
|
||
steps = []
|
||
for t in range(STEPS):
|
||
if t < ACTIVE:
|
||
p = kf.patches_at(r[t].astype(np.float32), pos, pos).astype(np.int8)
|
||
pot += (p.reshape(-1, p.shape[-1]).astype(np.float32) @ wq.astype(np.float32)).astype(np.int32).reshape(pot.shape) + bq
|
||
fired = pot >= thr1
|
||
cnt += fired
|
||
pot[fired] = 0
|
||
np.maximum(pot, 0, out=pot)
|
||
if per_step:
|
||
steps.append(pool(fired, edges))
|
||
res.append(np.stack(steps, axis=1) if per_step else pool(cnt, edges))
|
||
total.append(cnt.reshape(n, -1).sum(1))
|
||
return np.concatenate(res), np.concatenate(total)
|
||
|
||
|
||
def pool(c, edges):
|
||
c = c.reshape(c.shape[0], SIDE, SIDE, -1)
|
||
return np.concatenate([c[:, edges[gy]:edges[gy + 1], edges[gx]:edges[gx + 1]].sum((1, 2))
|
||
for gy in range(2) for gx in range(2)], axis=1).astype(np.int32)
|
||
|
||
|
||
def readout_neuron(steps, wq, bq, thr):
|
||
"""Exact readout neuron: steps [n, STEPS, F] layer-1 spikes arrive with delay 1; integer weights wq [F] and a
|
||
bias input bq on every active step; fires at thr, resets to 0, never below 0. Returns 1 (dog) if it fired."""
|
||
pot = np.zeros(len(steps), np.int64)
|
||
fired_any = np.zeros(len(steps), bool)
|
||
for t in range(1, STEPS):
|
||
pot += steps[:, t - 1] @ wq + (bq if t - 1 < ACTIVE else 0)
|
||
f = pot >= thr
|
||
fired_any |= f
|
||
pot[f] = 0
|
||
np.maximum(pot, 0, out=pot)
|
||
return fired_any.astype(int)
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("--model", required=True)
|
||
ap.add_argument("--data", default="/home/int/run/cats-dogs/data-v3")
|
||
ap.add_argument("--alpha", type=float, default=0.25)
|
||
ap.add_argument("--thr1", type=float, default=0, help="layer-1 threshold in quantised units (0: auto)")
|
||
ap.add_argument("--train-limit", type=int, default=0)
|
||
ap.add_argument("--out", required=True)
|
||
ap.add_argument("--active", type=int, default=10, help="active (spiking) input steps; 5 silent steps follow")
|
||
a = ap.parse_args()
|
||
globals().update(ACTIVE=a.active, STEPS=a.active + 5)
|
||
model = dict(np.load(a.model))
|
||
w, b = fold(model, a.alpha)
|
||
wq, bq, s = quant_layer1(w, b)
|
||
tr, va = np.load(f"{a.data}/train.npz"), np.load(f"{a.data}/val.npz")
|
||
xt = tr["rgb"] if not a.train_limit else tr["rgb"][:a.train_limit]
|
||
yt = tr["labels"][:len(xt)].astype(int)
|
||
if not a.thr1:
|
||
# Typical positive per-step drive of a layer-1 neuron on training patches.
|
||
p = kf.patches_at(xt[:200].astype(np.float32) / 255, range(SIDE), range(SIDE)).reshape(-1, 108)
|
||
act = p @ w + b
|
||
a.thr1 = float(np.percentile(act[act > 0], 99) * s)
|
||
print("layer-1 threshold", a.thr1, "weight scale", round(s, 2), flush=True)
|
||
t0 = time.time()
|
||
ct, tt = layer1_counts(np.concatenate([xt, xt[:, :, :, ::-1]]), wq, bq, a.thr1)
|
||
sv, tv = layer1_counts(va["rgb"], wq, bq, a.thr1, per_step=True)
|
||
cv = sv.sum(1)
|
||
print("simulated", len(ct) + len(cv), "images in", round(time.time() - t0), "s; layer-1 spikes per image",
|
||
int(tv.mean()), flush=True)
|
||
yy = np.r_[yt, yt]
|
||
np.savez(a.out, ct=ct, sv=sv, yt=yy, yv=va["labels"], thr1=a.thr1, wq=wq, bq=bq)
|
||
f = lambda c: np.hstack([c.astype(np.float32), np.ones((len(c), 1), np.float32)])
|
||
res = {}
|
||
for lam in [1e5, 1e6, 1e7]:
|
||
sol = np.linalg.solve(f(ct).T @ f(ct) + lam * np.eye(ct.shape[1] + 1), f(ct).T @ (yy * 2 - 1.0))
|
||
acc = 100 * (((f(cv) @ sol) > 0) == va["labels"]).mean()
|
||
rq = np.rint(sol / np.abs(sol).max() * 255).astype(np.int64)
|
||
# The readout bias input fires on each of the ACTIVE steps, so its weight is spread over them.
|
||
bias = int(np.rint(rq[-1] / ACTIVE))
|
||
# Readout threshold relative to the typical size of its per-step input.
|
||
unit = float(np.median(np.abs(sv[:, :ACTIVE].reshape(-1, sv.shape[-1]) @ rq[:-1] + bias))) + 1
|
||
spk = {q: 100 * (readout_neuron(sv, rq[:-1], bias, q * unit) == va["labels"]).mean() for q in (2, 3, 4, 5, 6)}
|
||
res[lam] = {"linear": acc, "spiking": spk, "unit": unit}
|
||
print(f"lam {lam:g}: linear on exact layer-1 counts {acc:.1f}%; spiking readout by threshold "
|
||
+ ", ".join(f"{q}×: {v:.1f}%" for q, v in spk.items()) + f" (unit {unit:.0f})", flush=True)
|
||
Path(a.out).with_suffix(".json").write_text(json.dumps({"thr1": a.thr1, "acc": res}, indent=1))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|