pooling.py

"""
THE HEAD HAS BEEN MOST OF THE MODEL ALL ALONG.

Every script in this programme ends with a head that reads EVERY hidden
unit: 3,136 of them here, so 31,370 parameters. Measured against it, the
folded body is 1.9% of a five-layer model at 28x28 and 82.5% at 7x7 — and
the depth sweep that produced those numbers could not separate "deeper
suits a larger grid" from "layers are cheap at a larger grid", because
scaling channels inversely with area made a layer cost 36,865 values at
7x7 and 145 at 28x28.

A conventional convolutional network does not do this. It POOLS over
spatial positions before the head, so the head reads one number per channel
rather than one per channel per position. Here that would take the head
from 31,370 parameters to CHANNELS x 10 — 650 at 7x7, 170 at 14x14, 50 at
28x28. The body would become the model instead of a rounding error.

Nothing in this programme has ever used pooling. That is not a small
omission: it is most of why these models are head-dominated, most of why
the storage ratios were inflated, and possibly why depth has looked so
weak.

So the same depth sweep is run twice, with and without pooling before the
head, and three things are read from it.

  DOES POOLING COST ACCURACY?   it discards where a feature fired and
                                keeps only how much, which on a centred
                                dataset may be free and may not be
  DOES IT CHANGE THE DEPTH STORY?  if depth only ever looked useful
                                because layers were cheap relative to a
                                huge head, pooling should change which
                                depth wins
  WHAT DOES IT COST TO STORE?   the head shrinks by a factor of the grid
                                area, which is 49 to 784 times

The honest expectation is that pooling costs some accuracy on a task this
small, and that the question is whether it costs less than the storage it
saves.
"""

import numpy as np
import time
import json
import os

try:
    import cupy as _cp
    _GPU = _cp.cuda.runtime.getDeviceCount() > 0
except Exception:
    _GPU = False
xp = _cp if _GPU else np
DT = np.float32


def to_dev(a, dtype=DT):
    a = np.asarray(a, dtype=dtype)
    return xp.asarray(a) if _GPU else a


def to_host(a):
    return _cp.asnumpy(a) if _GPU and isinstance(a, _cp.ndarray) else np.asarray(a)


def windowed(g, c_in, k, c_out):
    ni, no = c_in*g*g, c_out*g*g
    ii, jj = np.meshgrid(np.arange(ni), np.arange(no), indexing='ij')
    ci, pi = ii // (g*g), ii % (g*g)
    co, po = jj // (g*g), jj % (g*g)
    dr = pi // g - (po // g - k//2)
    dc = pi % g - (po % g - k//2)
    inside = (dr >= 0) & (dr < k) & (dc >= 0) & (dc < k)
    K = c_in*c_out*k*k + 1
    idx = np.where(inside, (ci*c_out + co)*k*k + dr*k + dc, K-1)
    return idx.ravel().astype(np.int32), K, no, no*(c_in*k*k)


DETERMINISTIC = True
_FIXED = {}


class FixedScatter:
    """Fixed-order accumulation, so a rerun reproduces exactly."""

    def __init__(self, idx, K, cap=8192):
        h = to_host(idx).astype(np.int64).reshape(-1)
        order = np.argsort(h, kind="stable")
        counts = np.bincount(h, minlength=K)
        starts = np.cumsum(counts) - counts
        big = np.where(counts > cap)[0]
        small = np.where(counts <= cap)[0]
        self.K = K
        self.order = to_dev(order, np.int64) if _GPU else order
        self.big = [(int(b), int(starts[b]), int(starts[b]+counts[b]))
                    for b in big]
        self.small = to_dev(small, np.int64) if _GPU else small
        self.width = int(counts[small].max()) if len(small) else 0
        if self.width:
            pos = np.concatenate([np.arange(counts[s]) for s in small])
            src = np.concatenate([np.arange(starts[s], starts[s]+counts[s])
                                  for s in small])
            row = np.repeat(np.arange(len(small)), counts[small])
            self.src = to_dev(src, np.int64) if _GPU else src
            sl = row*self.width + pos
            self.slot = to_dev(sl, np.int64) if _GPU else sl
            self.buf = xp.zeros(len(small)*self.width, DT)
        self._keep = idx

    def __call__(self, g):
        gs = g.reshape(-1)[self.order]
        out = xp.zeros(self.K, DT)
        if self.width:
            self.buf[:] = 0
            self.buf[self.slot] = gs[self.src]
            out[self.small] = self.buf.reshape(-1, self.width).sum(1)
        for b, a, z in self.big:
            out[b] = gs[a:z].sum()
        return out


def scatter(dW, idx, K):
    if not DETERMINISTIC:
        g = xp.zeros(K, DT)
        if _GPU:
            import cupyx
            cupyx.scatter_add(g, idx, dW.reshape(-1))
        else:
            np.add.at(g, idx, dW.reshape(-1))
        return g
    key = (id(idx), K)
    if key not in _FIXED:
        _FIXED[key] = FixedScatter(idx, K)
    return _FIXED[key](dW)


def train(Xtr, Ytr, Xte, yte, g, chan, depth, cfg, seed, pool=False):
    D = Xtr.shape[1]
    rg = np.random.default_rng(seed)
    layers, cin = [], 1
    for l in range(depth):
        idx, K, no, macs = windowed(g, cin, 3, chan)
        layers.append(dict(idx=to_dev(idx, np.int32) if _GPU else idx,
                           K=K, out=no, macs=macs,
                           ins=D if l == 0 else layers[-1]["out"],
                           taps=cin*9))
        cin = chan
    L = depth
    P = []
    for l in layers:
        v = rg.normal(0, np.sqrt(2.0/l["taps"]), l["K"]).astype(np.float32)
        v[-1] = 0.0
        P.append(to_dev(v))
    P += [xp.ones(l["out"], DT) for l in layers]
    P += [xp.zeros(l["out"], DT) for l in layers]
    hw = chan if pool else layers[-1]["out"]
    P += [to_dev(rg.normal(0, np.sqrt(2.0/hw), (hw, 10))), xp.zeros(10, DT)]
    HEAD, OB = 3*L, 3*L+1
    M = [xp.zeros_like(p) for p in P]; V = [xp.zeros_like(p) for p in P]
    n = Xtr.shape[0]; t = 0
    ag = np.random.default_rng(seed + 991)

    npos = g*g

    def fwd(x, keep=False):
        cache = []; h = x
        for li, l in enumerate(layers):
            W = P[li][l["idx"]].reshape(l["ins"], l["out"])
            z = h @ W
            var = z.var(1, keepdims=True) + 1e-5
            zn = (z - z.mean(1, keepdims=True))/xp.sqrt(var)
            zs = zn*P[L+li] + P[2*L+li]
            a = xp.maximum(zs, 0)
            if keep:
                cache.append((h, W, var, zn, zs))
            h = a
        if pool:
            # average over spatial positions: the head then reads one
            # number per channel instead of one per channel per position
            h = h.reshape(h.shape[0], chan, npos).mean(2)
        return h, cache

    for ep in range(cfg["epochs"]):
        perm = ag.permutation(n)
        for st in range(0, n, cfg["batch"]):
            b = perm[st:st+cfg["batch"]]
            x = Xtr[b]; y = Ytr[b]
            h, cache = fwd(x, keep=True)
            lg = h @ P[HEAD] + P[OB]
            e = xp.exp(lg - lg.max(1, keepdims=True))
            d = (e/e.sum(1, keepdims=True) - y)/len(b)
            G = [xp.zeros_like(p) for p in P]
            G[HEAD] = h.T @ d; G[OB] = d.sum(0)
            dh = d @ P[HEAD].T
            if pool:
                dh = xp.broadcast_to(dh[:, :, None]/npos,
                                     (dh.shape[0], chan, npos)
                                     ).reshape(dh.shape[0], chan*npos)
            for li in range(L-1, -1, -1):
                hin, W, var, zn, zs = cache[li]
                dzs = dh*(zs > 0)
                G[L+li] = (dzs*zn).sum(0); G[2*L+li] = dzs.sum(0)
                dzn = dzs*P[L+li]
                dz = (dzn - dzn.mean(1, keepdims=True)
                      - zn*(dzn*zn).mean(1, keepdims=True))/xp.sqrt(var)
                G[li] = scatter(hin.T @ dz, layers[li]["idx"], layers[li]["K"])
                if li > 0:
                    dh = dz @ W.T
            t += 1
            for i, (p_, gr) in enumerate(zip(P, G)):
                M[i] = 0.9*M[i] + 0.1*gr
                V[i] = 0.999*V[i] + 0.001*gr*gr
                P[i] = p_ - cfg["lr"]*(M[i]/(1-0.9**t)) \
                    / (xp.sqrt(V[i]/(1-0.999**t))+1e-8)
    out = []
    for s in range(0, Xte.shape[0], 4096):
        h, _ = fwd(Xte[s:s+4096])
        out.append(to_host(h @ P[HEAD] + P[OB]))
    acc = float((np.concatenate(out).argmax(1) == yte).mean())
    # COUNT THE HEAD. It is dense and unfolded, and on Fashion at 14x14 it
    # is 99.5% of a one-layer model's parameters — reporting the folded
    # body alone made a 15% storage difference look like 65x.
    head = (chan if pool else layers[-1]["out"])*10 + 10
    return (acc, sum(l["K"] for l in layers) + head,
            sum(l["macs"] for l in layers) + head - 10, head)


def load(grid, cfg):
    from tensorflow import keras
    (a, b), (c, d) = keras.datasets.fashion_mnist.load_data()
    X = np.concatenate([a, c]).astype(np.float32)/255.0
    y = np.concatenate([b, d]).ravel().astype(np.int64)
    if grid != 28:
        s = 28//grid
        X = X.reshape(-1, grid, s, grid, s).mean(axis=(2, 4))
    rg = np.random.default_rng(0); p = rg.permutation(len(X))
    tr, te = p[:cfg["n_train"]], p[cfg["n_train"]:cfg["n_train"]+10000]
    mu, sd = X[tr].mean(), X[tr].std()+1e-8
    f = lambda Z: ((Z-mu)/sd).reshape(len(Z), -1)
    Y = np.zeros((len(tr), 10), np.float32); Y[np.arange(len(tr)), y[tr]] = 1
    return f(X[tr]), Y, f(X[te]), y[te]


CFG = dict(n_train=20000, batch=128, lr=1e-3, epochs=40, seeds=(0, 1, 2),
           hidden_target=3136, grids=(7, 14, 28), depths=(1, 2, 3, 4, 5))


def main(**over):
    CFG.update(over)
    t0 = time.time()
    print("=" * 78)
    print("THE HEAD HAS BEEN MOST OF THE MODEL ALL ALONG")
    print("=" * 78)
    print(f"  backend: {'cupy (GPU)' if _GPU else 'numpy (CPU)'}")
    for k, v in CFG.items():
        print(f"  {k:14s} = {v}")
    print(f"\n  channels scale inversely with area to hold the hidden width")
    print(f"  near {CFG['hidden_target']}, so no arm is favoured by having a")
    print(f"  bigger matrix product:")
    for g in CFG["grids"]:
        c = max(1, CFG["hidden_target"]//(g*g))
        print(f"    {g:2d}x{g:<2d} -> {c:3d} channels, hidden {c*g*g:,}")
    print(f"\n  reach says L >= (n-1)/2 to cover the grid: "
          + ", ".join(f"{g}x{g} needs {int(np.ceil((g-1)/2))}"
                      for g in CFG["grids"]))
    print("=" * 78, flush=True)

    res = {}
    for g in CFG["grids"]:
        chan = max(1, CFG["hidden_target"]//(g*g))
        Xtr, Ytr, Xte, yte = load(g, CFG)
        Xtr, Ytr, Xte = to_dev(Xtr), to_dev(Ytr), to_dev(Xte)
        print(f"\n  {g}x{g}, {chan} channels")
        print(f"  (the head alone is "
              f"{chan*g*g*10 + 10:,} parameters, counted in every row)")
        print(f"  {'head':>7s} {'L':>4s} {'TOTAL vals':>11s} "
              f"{'multiplies':>12s} {'accuracy':>9s} {'sd':>7s} {'reach':>7s}")
        for pool, L in [(p, L) for p in (False, True)
                        for L in CFG["depths"]]:
            # the fixed-order tables are keyed by the index array's identity
            # and each arm builds new ones, so the cache would grow across
            # fifteen arms until the device ran out
            _FIXED.clear()
            if _GPU:
                _cp.get_default_memory_pool().free_all_blocks()
            out = [train(Xtr, Ytr, Xte, yte, g, chan, L, CFG, s,
                         pool=pool)
                   for s in CFG["seeds"]]
            acc = [o[0] for o in out]
            K, mc, hd = out[0][1], out[0][2], out[0][3]
            reach = min(g, 2*L+1)
            res[f"{g}/{L}/{int(pool)}"] = dict(acc=float(np.mean(acc)),
                                   sd=float(np.std(acc)), K=K, macs=mc,
                                   reach=reach, chan=chan, head=hd)
            print(f"  {'pooled' if pool else 'flat':>7s} {L:4d} "
                  f"{K:11,} {mc:12,} {np.mean(acc):9.4f} "
                  f"{np.std(acc):7.4f} {reach:5d}x{reach:<2d}"
                  f"   [{time.time()-t0:.0f}s]", flush=True)
            json.dump(res, open("pooling.json", "w"), indent=2)

    print("\n" + "=" * 78)
    print("  WHAT DOES POOLING COST, AND WHAT DOES IT SAVE?")
    print("=" * 78)
    print(f"  {'grid':>5s} {'L':>3s} {'flat':>8s} {'pooled':>8s} "
          f"{'accuracy':>9s} {'head flat':>10s} {'head pooled':>12s} "
          f"{'saved':>7s}")
    for g in CFG["grids"]:
        for L in CFG["depths"]:
            a = res.get(f"{g}/{L}/0"); b = res.get(f"{g}/{L}/1")
            if not (a and b):
                continue
            print(f"  {g:5d} {L:3d} {a['acc']:8.4f} {b['acc']:8.4f} "
                  f"{b['acc']-a['acc']:+9.4f} {a['head']:10,} "
                  f"{b['head']:12,} {a['head']/b['head']:6.0f}x")

    print("\n" + "=" * 78)
    print("  READOUT")
    print("=" * 78)
    sd = max(r["sd"] for r in res.values())
    d = [res[f"{g}/{L}/1"]["acc"] - res[f"{g}/{L}/0"]["acc"]
         for g in CFG["grids"] for L in CFG["depths"]
         if f"{g}/{L}/1" in res and f"{g}/{L}/0" in res]
    d = np.array(d)
    print(f"  pooling changes accuracy by {d.mean():+.4f} on average, "
          f"range {d.min():+.4f} to {d.max():+.4f}")
    print(f"  seed spread (worst) {sd:.4f}\n")
    print(f"  best depth, flat head vs pooled:")
    for g in CFG["grids"]:
        bf = max(CFG["depths"], key=lambda L: res[f"{g}/{L}/0"]["acc"])
        bp = max(CFG["depths"], key=lambda L: res[f"{g}/{L}/1"]["acc"])
        print(f"    {g:2d}x{g:<2d}  flat {bf}   pooled {bp}"
              + ("   <- the depth story CHANGES" if bf != bp else ""))
    print()
    if d.mean() > -2*sd:
        print("  POOLING IS FREE. The head shrinks by the grid area and")
        print("  accuracy does not move, so every model in this programme")
        print("  has been carrying a head one to two orders of magnitude")
        print("  larger than it needed — which is most of why the storage")
        print("  ratios were inflated and the body looked negligible.")
    elif d.mean() > -0.02:
        print(f"  POOLING COSTS {abs(d.mean()):.4f} AND SAVES THE HEAD. Whether")
        print("  that is worth it depends on the budget, but it should be a")
        print("  stated choice rather than an omission — and the ratios in")
        print("  the paper were computed without it.")
    else:
        print(f"  POOLING COSTS TOO MUCH ({d.mean():+.4f}). Discarding WHERE a")
        print("  feature fired matters on this task, so the large head is")
        print("  earning its storage and the programme's design was right by")
        print("  accident rather than wrong.")
    print(f"\n  total {time.time()-t0:.0f}s; wrote pooling.json")


if __name__ == "__main__":
    main()