118 lines
4.6 KiB
Python
118 lines
4.6 KiB
Python
"""Offline simulator of the quantized AltAI mnist-learn network at inference.
|
|
|
|
Reproduces KNP inference: raster -> INPUT (delay 1), optional INPUT -> INPUT inhibition (delay 1),
|
|
INPUT -> OUTPUT column vote (delay 1), 15 steps per image, potentials carried across images.
|
|
|
|
usage: python3 altai_sim.py <dump_dir> <images.bin> <labels.target> [--inhib 255] [--thr 40,60,80] [--first N] [--test 10000]
|
|
[--pix 1568] (ON/OFF input, KNP_MNIST_INPUT_SIZE)
|
|
"""
|
|
import argparse
|
|
import csv
|
|
import numpy as np
|
|
|
|
STEPS, ACTIVE, CLASSES, PIX = 15, 10, 10, 784
|
|
|
|
|
|
def load_dump(d, tag="quant"):
|
|
with open(f"{d}/{tag}_neurons.csv") as f:
|
|
rows = list(csv.DictReader(f))
|
|
thr = np.array([float(r["activation_threshold"]) + float(r["additional_threshold"]) for r in rows])
|
|
w = np.zeros((PIX, len(rows)))
|
|
with open(f"{d}/{tag}_weights.csv") as f:
|
|
next(f)
|
|
for line in f:
|
|
pre, post, wt = line.split(",")
|
|
w[int(pre), int(post)] = float(wt)
|
|
return w, thr
|
|
|
|
|
|
def load_images(path, label_path, first, count):
|
|
imgs = np.fromfile(path, dtype=np.uint8).reshape(-1, PIX)[first:first + count]
|
|
labels = np.array([int(x) for x in open(label_path).read().split()])[first:first + count]
|
|
return imgs, labels
|
|
|
|
|
|
def raster(img, states):
|
|
"""Incrementing image-to-spikes converter; `states` persists between images like in KNP."""
|
|
out = np.zeros((STEPS, PIX), dtype=bool)
|
|
inc = img.astype(np.float32) * np.float32(1 / 255.0)
|
|
for t in range(ACTIVE):
|
|
states += inc
|
|
fired = states >= 1.0
|
|
out[t] = fired
|
|
states[fired] -= 1.0
|
|
return out
|
|
|
|
|
|
class Sim:
|
|
def __init__(self, w, thr, inhib=0.0):
|
|
self.w, self.thr, self.inhib = w, thr, inhib
|
|
n = len(thr)
|
|
self.col = n // CLASSES
|
|
self.pot = np.zeros(n)
|
|
self.states = np.zeros(PIX, dtype=np.float32)
|
|
self.prev_raster = np.zeros(PIX, dtype=bool)
|
|
self.prev_input = np.zeros(n, dtype=bool)
|
|
|
|
def run_image(self, img, trace=False):
|
|
"""Returns (votes per class, per-step trace if requested)."""
|
|
spikes = raster(img, self.states)
|
|
votes = np.zeros(CLASSES, dtype=int)
|
|
steps = []
|
|
for t in range(STEPS):
|
|
# OUTPUT: column fires if any of its INPUT neurons fired on the previous step.
|
|
col_fired = self.prev_input.reshape(CLASSES, self.col).any(axis=1)
|
|
votes += col_fired
|
|
# INPUT: impacts from previous-step raster and inhibition.
|
|
self.pot = np.round(self.pot)
|
|
self.pot += self.prev_raster @ self.w
|
|
if self.inhib and self.prev_input.any():
|
|
k = self.prev_input.sum()
|
|
self.pot -= self.inhib * (k - self.prev_input) # no self-inhibition
|
|
fired = self.pot >= self.thr
|
|
self.pot[fired] = 0
|
|
self.pot = np.maximum(self.pot, 0) # negative threshold 0 with saturation
|
|
if trace:
|
|
steps.append({"raster": spikes[t].copy(), "pot": self.pot.copy(), "fired": fired.copy(),
|
|
"votes": votes.copy()})
|
|
self.prev_raster, self.prev_input = spikes[t], fired
|
|
return votes, steps
|
|
|
|
|
|
def predict(votes):
|
|
return int(np.argmax(votes)) if votes.max() > 0 else -1
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("dump")
|
|
ap.add_argument("images")
|
|
ap.add_argument("labels")
|
|
ap.add_argument("--inhib", type=float, default=0)
|
|
ap.add_argument("--first", type=int, default=None, help="index of first test image (default: last --test)")
|
|
ap.add_argument("--test", type=int, default=10000)
|
|
ap.add_argument("--thr", default="", help="comma-separated thresholds to sweep instead of the dumped ones")
|
|
ap.add_argument("--pix", type=int, default=784, help="input size per image (1568 for ON/OFF)")
|
|
a = ap.parse_args()
|
|
global PIX
|
|
PIX = a.pix
|
|
w, thr = load_dump(a.dump)
|
|
total = len(open(a.labels).read().split())
|
|
first = total - a.test if a.first is None else a.first
|
|
imgs, labels = load_images(a.images, a.labels, first, a.test)
|
|
for t in [float(x) for x in a.thr.split(",")] if a.thr else [None]:
|
|
sim = Sim(w, thr if t is None else np.full_like(thr, t), a.inhib)
|
|
tp = fn = fp = 0
|
|
for img, lab in zip(imgs, labels):
|
|
p = predict(sim.run_image(img)[0])
|
|
if p < 0:
|
|
fn += 1
|
|
elif p == lab:
|
|
tp += 1
|
|
else:
|
|
fp += 1
|
|
tag = "" if t is None else f"thr {t:g} "
|
|
print(f"{tag}correct {tp} unanswered {fn} wrong {fp} accuracy {tp / len(labels):.4f}")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|