79 lines
3.8 KiB
Python
79 lines
3.8 KiB
Python
"""Train the class neurons offline on detector spike counts produced by KNP (altai_infer), write integer weights.
|
|
|
|
Ridge regression (one-vs-rest, targets ±1) on per-(quadrant, detector type) spike counts of the training images
|
|
(+ mirrored copies). Standardisation is folded into the weights and bias; weights are scaled to integers in ±255,
|
|
the bias is spread over the active steps. The class-neuron threshold is chosen on the training counts with a rate
|
|
approximation (spikes ≈ positive drive / threshold) so that the most-spiking class matches the ridge decision.
|
|
usage: ~/anvenv/bin/python train_readout.py <model_dir> <counts_dir> [--lam 1000]
|
|
"""
|
|
import argparse
|
|
import re
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
ACTIVE = 10
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("model")
|
|
ap.add_argument("counts")
|
|
ap.add_argument("--lam", type=float, default=1000)
|
|
ap.add_argument("--data", default=str(Path.home() / "run/cifar/vehicles"))
|
|
a = ap.parse_args()
|
|
model, cdir = Path(a.model), Path(a.counts)
|
|
k, stride, det_thr, classes, _ = (int(v) for v in (model / "meta.txt").read_text().split())
|
|
|
|
|
|
def load(tag, n):
|
|
pattern = re.compile(re.escape(tag) + r"_(\d+)_counts\.bin")
|
|
parts = sorted((p for p in cdir.glob('*_counts.bin') if pattern.fullmatch(p.name)),
|
|
key=lambda p: int(pattern.fullmatch(p.name).group(1)))
|
|
arrays, next_index = [], 0
|
|
for path in parts:
|
|
first = int(pattern.fullmatch(path.name).group(1))
|
|
if first != next_index or not (cdir / f'{tag}_{first}.done').exists():
|
|
raise ValueError(f'{tag}: missing checkpoint or noncontiguous range at {first}')
|
|
raw = np.fromfile(path, np.int32)
|
|
if len(raw) % (4 * k):
|
|
raise ValueError(f'Truncated counts: {path}')
|
|
block = raw.reshape(-1, 4 * k)
|
|
arrays.append(block)
|
|
next_index += len(block)
|
|
if next_index != n:
|
|
raise ValueError(f'{tag}: expected {n} rows, got {next_index}')
|
|
return np.concatenate(arrays).astype(np.float32)
|
|
|
|
|
|
yt = np.array(open(f"{a.data}/train_labels.txt").read().split(), int)
|
|
yv = np.array(open(f"{a.data}/val_labels.txt").read().split(), int)
|
|
xt = np.vstack([load("train", len(yt)), load("train_flip", len(yt))])
|
|
yt2 = np.r_[yt, yt]
|
|
xv = load("val", len(yv))
|
|
mu, sd = xt.mean(0), xt.std(0) + 1e-3
|
|
f = lambda z: np.hstack([(z - mu) / sd, np.ones((len(z), 1), np.float32)])
|
|
T = np.eye(classes)[yt2] * 2 - 1
|
|
for lam in sorted({a.lam, 100.0, 10000.0}):
|
|
sol = np.linalg.solve(f(xt).T @ f(xt) + lam * np.eye(4 * k + 1), f(xt).T @ T)
|
|
print(f"lam {lam:g}: ridge on KNP counts, validation {100 * ((f(xv) @ sol).argmax(1) == yv).mean():.2f}%")
|
|
sol = np.linalg.solve(f(xt).T @ f(xt) + a.lam * np.eye(4 * k + 1), f(xt).T @ T)
|
|
w = sol[:-1] / sd[:, None] # on raw counts
|
|
b = sol[-1] - (mu / sd) @ sol[:-1]
|
|
s = 255 / np.abs(w).max()
|
|
wq = np.rint(w * s).astype(np.int32) # [4K, classes]
|
|
bq_total = b * s
|
|
bq = np.rint(bq_total / ACTIVE).astype(np.int32)
|
|
drive_t = xt @ wq + bq * ACTIVE # total drive per image and class
|
|
best = None
|
|
# A class neuron fires at most once a step (about ACTIVE + 1 spikes per image), so the threshold is a fraction
|
|
# of the typical positive per-step drive; spikes are capped accordingly.
|
|
unit = float(np.percentile(drive_t[drive_t > 0], 50)) / ACTIVE
|
|
for q in (0.1, 0.2, 0.3, 0.5, 0.8):
|
|
thr = int(min(max(unit * q, 1), 65535))
|
|
spikes = np.minimum(np.floor(np.maximum(drive_t, 0) / thr), ACTIVE + 1)
|
|
acc = 100 * (spikes.argmax(1) == yt2).mean()
|
|
print(f"threshold {thr}: train accuracy (rate approximation) {acc:.2f}%, max spikes {spikes.max():.0f}")
|
|
if best is None or acc > best[0]:
|
|
best = (acc, thr)
|
|
wq.T.copy().tofile(model / "readout_w.bin")
|
|
bq.tofile(model / "readout_b.bin")
|
|
(model / "meta.txt").write_text(f"{k} {stride} {det_thr} {classes} {best[1]}\n")
|
|
print("written", model, "threshold", best[1])
|