92 lines
4.3 KiB
Python
92 lines
4.3 KiB
Python
"""Learned local features for small colour images (Coates & Ng 2011 style), prototype for a spiking layer.
|
||
|
||
Layer 1 is learned without labels: random 6×6×3 patches are contrast-normalised and ZCA-whitened, then
|
||
spherical k-means finds K prototypes. This is the same rule as winner-take-all competitive Hebbian learning:
|
||
the closest prototype moves towards the patch. Activation of prototype k at a position is
|
||
max(0, <patch, d_k> - alpha) (a neuron with threshold alpha). Activations are sum-pooled over a grid of
|
||
regions, which gives tolerance to position.
|
||
|
||
API: learn(images) -> model dict; encode(images, model) -> float32 features [n, K * grid * grid].
|
||
images: uint8 [n, 3, 32, 32].
|
||
"""
|
||
import numpy as np
|
||
|
||
PATCH = 6
|
||
|
||
|
||
def patches_at(img, ys, xs):
|
||
"""img [n,3,32,32] float32 → patches [n, len(ys)*len(xs), 108]."""
|
||
out = np.stack([img[:, :, y:y + PATCH, x:x + PATCH].reshape(len(img), -1) for y in ys for x in xs], axis=1)
|
||
return out
|
||
|
||
|
||
# CONTRAST_DIVIDE=False keeps only the (linear) patch mean subtraction, so mean removal, ZCA and the
|
||
# prototypes fold into one weight matrix of a spiking neuron; division by patch contrast is not linear.
|
||
CONTRAST_DIVIDE = True
|
||
|
||
|
||
def normalise(p):
|
||
p = p - p.mean(-1, keepdims=True)
|
||
return p / np.sqrt(p.var(-1, keepdims=True) + 10 / 255) if CONTRAST_DIVIDE else p * 4
|
||
|
||
|
||
def learn(images, k=400, n_patches=200000, iters=15, seed=0):
|
||
rng = np.random.default_rng(seed)
|
||
x = images.astype(np.float32) / 255
|
||
idx = rng.integers(0, len(x), n_patches)
|
||
ys, xs = rng.integers(0, 32 - PATCH + 1, n_patches), rng.integers(0, 32 - PATCH + 1, n_patches)
|
||
p = np.stack([x[i, :, y:y + PATCH, xx:xx + PATCH].reshape(-1) for i, y, xx in zip(idx, ys, xs)])
|
||
p = normalise(p)
|
||
mean = p.mean(0)
|
||
cov = np.cov(p - mean, rowvar=False)
|
||
e, v = np.linalg.eigh(cov)
|
||
zca = (v * (1 / np.sqrt(e + 0.1))) @ v.T
|
||
pw = (p - mean) @ zca
|
||
d = rng.normal(size=(k, pw.shape[1])).astype(np.float32)
|
||
d /= np.linalg.norm(d, axis=1, keepdims=True)
|
||
for _ in range(iters):
|
||
s = pw @ d.T
|
||
win = s.argmax(1)
|
||
new = np.zeros_like(d)
|
||
np.add.at(new, win, pw * s[np.arange(len(pw)), win, None])
|
||
empty = np.linalg.norm(new, axis=1) < 1e-6
|
||
new[empty] = rng.normal(size=(empty.sum(), d.shape[1]))
|
||
d = new / np.linalg.norm(new, axis=1, keepdims=True)
|
||
return {"mean": mean.astype(np.float32), "zca": zca.astype(np.float32), "d": d.astype(np.float32)}
|
||
|
||
|
||
def encode(images, model, alpha=0.25, grid=2, batch=250):
|
||
x = images.astype(np.float32) / 255
|
||
pos = list(range(0, 32 - PATCH + 1))
|
||
side = len(pos)
|
||
edges = np.linspace(0, side, grid + 1).astype(int)
|
||
out = []
|
||
for i in range(0, len(x), batch):
|
||
p = normalise(patches_at(x[i:i + batch], pos, pos))
|
||
a = np.maximum(((p - model["mean"]) @ model["zca"]) @ model["d"].T - alpha, 0)
|
||
a = a.reshape(len(p), side, side, -1)
|
||
pooled = [a[:, edges[gy]:edges[gy + 1], edges[gx]:edges[gx + 1]].sum((1, 2))
|
||
for gy in range(grid) for gx in range(grid)]
|
||
out.append(np.concatenate(pooled, axis=1))
|
||
return np.concatenate(out)
|
||
|
||
|
||
def learn_online(images, k=400, n_patches=200000, lr0=0.05, seed=0, base=None):
|
||
"""Online winner-take-all learning, one patch at a time (what a spiking WTA layer with Hebbian/STDP does):
|
||
the most responsive neuron moves its normalised weights towards the patch, the others stay. Uses the same
|
||
whitening as `learn` (from `base` if given). Single pass, learning rate decays as lr0 / (1 + t / 20000)."""
|
||
rng = np.random.default_rng(seed)
|
||
if base is None:
|
||
base = learn(images, k=k, n_patches=min(n_patches, 50000), iters=0, seed=seed)
|
||
x = images.astype(np.float32) / 255
|
||
idx = rng.integers(0, len(x), n_patches)
|
||
ys, xs = rng.integers(0, 32 - PATCH + 1, n_patches), rng.integers(0, 32 - PATCH + 1, n_patches)
|
||
p = np.stack([x[i, :, y:y + PATCH, xx:xx + PATCH].reshape(-1) for i, y, xx in zip(idx, ys, xs)])
|
||
pw = (normalise(p) - base["mean"]) @ base["zca"]
|
||
d = pw[rng.choice(len(pw), k, replace=False)].copy()
|
||
d /= np.linalg.norm(d, axis=1, keepdims=True)
|
||
for t, v in enumerate(pw):
|
||
win = int(np.argmax(d @ v))
|
||
d[win] += lr0 / (1 + t / 20000) * v
|
||
d[win] /= np.linalg.norm(d[win])
|
||
return {"mean": base["mean"], "zca": base["zca"], "d": d.astype(np.float32)}
|