113 lines
4.7 KiB
Python
113 lines
4.7 KiB
Python
"""Design check for the all-on-platform pipeline (numpy prototype, not the final result).
|
||
|
||
Retina: per colour channel a centre–surround (difference of Gaussians) filter; ON = max(0, DoG), OFF = max(0, −DoG)
|
||
→ 6 maps 32×32 (fixed weights, i.e. one fixed layer of neurons). Features: online winner-take-all on 6×6×6 retina
|
||
patches (unit-norm weights, winner = largest w·x, winner moves towards the patch), activation max(0, w·x − θ) with
|
||
θ a quantile of the winners' responses; sum pooling over quadrants; linear readout (ridge here, delta rule in KNP).
|
||
|
||
usage: ~/anvenv/bin/python retina_proto.py <data dir with train.npz/val.npz> [--k 200 --stride 2]
|
||
"""
|
||
import argparse
|
||
|
||
import numpy as np
|
||
|
||
PATCH = 6
|
||
|
||
|
||
def gauss(sigma, r=3):
|
||
y, x = np.mgrid[-r:r + 1, -r:r + 1]
|
||
g = np.exp(-(x * x + y * y) / (2 * sigma * sigma))
|
||
return g / g.sum()
|
||
|
||
|
||
DOG = gauss(0.8) - gauss(2.0)
|
||
|
||
|
||
def retina(rgb, gain=4.0):
|
||
"""rgb uint8 [n,3,32,32] → float32 [n,6,32,32] in [0,1]: ON and OFF maps per colour."""
|
||
x = rgb.astype(np.float32) / 255
|
||
r = DOG.shape[0] // 2
|
||
xp = np.pad(x, ((0, 0), (0, 0), (r, r), (r, r)), mode="edge")
|
||
d = np.zeros_like(x)
|
||
for dy in range(DOG.shape[0]):
|
||
for dx in range(DOG.shape[1]):
|
||
d += DOG[dy, dx] * xp[:, :, dy:dy + 32, dx:dx + 32]
|
||
return np.clip(np.concatenate([np.maximum(d, 0), np.maximum(-d, 0)], axis=1) * gain, 0, 1)
|
||
|
||
|
||
def patches(maps, stride):
|
||
pos = range(0, 32 - PATCH + 1, stride)
|
||
return np.stack([maps[:, :, y:y + PATCH, x:x + PATCH].reshape(len(maps), -1) for y in pos for x in pos], axis=1)
|
||
|
||
|
||
def learn_wta(maps, k, n=300000, lr0=0.05, seed=0):
|
||
rng = np.random.default_rng(seed)
|
||
i, y, x = rng.integers(0, len(maps), n), rng.integers(0, 27, n), rng.integers(0, 27, n)
|
||
p = np.stack([maps[a, :, b:b + PATCH, c:c + PATCH].reshape(-1) for a, b, c in zip(i, y, x)])
|
||
p = p[p.sum(1) > 0.5] # skip empty patches (no retina activity)
|
||
w = p[rng.choice(len(p), k, replace=False)].copy()
|
||
w /= np.linalg.norm(w, axis=1, keepdims=True)
|
||
for t, v in enumerate(p):
|
||
win = int(np.argmax(w @ v))
|
||
w[win] += lr0 / (1 + t / 20000) * (v / (np.linalg.norm(v) + 1e-6) - w[win])
|
||
w[win] /= np.linalg.norm(w[win])
|
||
return w, p
|
||
|
||
|
||
def encode(maps, w, theta, stride, batch=500, centre=False):
|
||
"""centre=True: subtract the mean prototype (activation max(0, (w_k − mean_j w_j)·x − θ)); still linear."""
|
||
if centre:
|
||
w = w - w.mean(0)
|
||
out = []
|
||
for s in range(0, len(maps), batch):
|
||
p = patches(maps[s:s + batch], stride)
|
||
a = np.maximum(p @ w.T - theta, 0)
|
||
side = int(np.sqrt(p.shape[1]))
|
||
a = a.reshape(len(p), side, side, -1)
|
||
h = side // 2
|
||
out.append(np.concatenate([a[:, :h, :h].sum((1, 2)), a[:, :h, h:].sum((1, 2)),
|
||
a[:, h:, :h].sum((1, 2)), a[:, h:, h:].sum((1, 2))], axis=1))
|
||
return np.concatenate(out)
|
||
|
||
|
||
def ridge_acc(ft, yt, fv, yv, lam, classes):
|
||
mu, sd = ft.mean(0), ft.std(0) + 1e-3
|
||
f = lambda z: np.hstack([(z - mu) / sd, np.ones((len(z), 1), np.float32)])
|
||
w = np.linalg.solve(f(ft).T @ f(ft) + lam * np.eye(ft.shape[1] + 1), f(ft).T @ (np.eye(classes)[yt] * 2 - 1))
|
||
return 100 * ((f(fv) @ w).argmax(1) == yv).mean()
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("data")
|
||
ap.add_argument("--k", type=int, default=200)
|
||
ap.add_argument("--stride", type=int, default=2)
|
||
ap.add_argument("--flip", action="store_true")
|
||
ap.add_argument("--centre", action="store_true")
|
||
ap.add_argument("--sigmas", default="0.8,2.0")
|
||
ap.add_argument("--gain", type=float, default=4.0)
|
||
a = ap.parse_args()
|
||
global DOG
|
||
s1, s2 = (float(v) for v in a.sigmas.split(","))
|
||
DOG = gauss(s1) - gauss(s2)
|
||
tr, va = np.load(f"{a.data}/train.npz"), np.load(f"{a.data}/val.npz")
|
||
xt, yt, xv, yv = tr["rgb"], tr["labels"].astype(int), va["rgb"], va["labels"].astype(int)
|
||
classes = int(yt.max()) + 1
|
||
mt, mv = retina(xt, a.gain), retina(xv, a.gain)
|
||
w, p = learn_wta(mt, a.k)
|
||
wc = w - w.mean(0) if a.centre else w
|
||
resp = (p[:20000] @ wc.T).max(1)
|
||
for q in ((0, 50) if a.centre else (50, 75)):
|
||
theta = float(np.percentile(resp, q)) if q else 0.0
|
||
ft, fv = encode(mt, w, theta, a.stride, centre=a.centre), encode(mv, w, theta, a.stride, centre=a.centre)
|
||
if a.flip:
|
||
ft = np.vstack([ft, encode(retina(xt[:, :, :, ::-1], a.gain), w, theta, a.stride, centre=a.centre)])
|
||
ytt = np.r_[yt, yt]
|
||
else:
|
||
ytt = yt
|
||
print(f"K {a.k} stride {a.stride} theta=q{q} ({theta:.3f}): "
|
||
+ " | ".join(f"lam {lam}: {ridge_acc(ft, ytt, fv, yv, lam, classes):.1f}%" for lam in (100, 1000)), flush=True)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|