grown.py

"""
GROWING A MODEL FROM ITS OWN PARENTS, MID-TRAINING.

Merging two trained folded models failed. The channel permutation was real
-- alignment moved 14.7 of 16 channels -- and repairing it recovered a
great deal, but the child still landed at 0.4023 against parents at 0.6210.
Two runs differ by more than a relabelling.

Concatenation avoids the problem entirely: give the child BOTH parents'
channels rather than asking them to share slots. Nothing is averaged and
nothing is lost. But concatenation with stacked heads is exactly prediction
averaging -- verified to one part in ten trillion -- so on its own it buys
nothing, at identical values and identical arithmetic.

What it opens is co-adaptation. Join the parents BEFORE they finish and the
halves train together afterwards: a channel in half A is now optimised in
the presence of half B, so it can stop covering what B already covers and
specialise into what B misses. Neither parent could do that alone, and two
finished models stapled together cannot either.

The child is initialised so that it STARTS at exactly its parents'
ensemble: the heads are stacked and scaled. Every point it gains after that
is co-adaptation, measured against the thing it began as.

The controls are what make this answerable, and one of them is obvious
enough to be fatal if omitted.

  WIDE FROM SCRATCH   the same final width, trained from epoch zero. If
                      growing does not beat this, the parents' early
                      learning was worth nothing and this is a long route
                      to a bigger model.
  ENSEMBLE            the parents trained to the end and averaged, which is
                      what the child started as and had to beat.
  PARENTS             each alone.

and the join point is swept, because "there is more to learn from" is a
claim about WHEN: joining early should beat joining late if co-adaptation
is the mechanism, and the two should be equal if it is not.
"""

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)


def scatter(dW, idx, K):
    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


def augment(Xb, side, pad, rg):
    n = Xb.shape[0]
    im = Xb.reshape(n, 3, side, side)
    if pad:
        P = xp.zeros((n, 3, side+2*pad, side+2*pad), dtype=Xb.dtype)
        P[:, :, pad:pad+side, pad:pad+side] = im
        oy = rg.integers(0, 2*pad+1, n); ox = rg.integers(0, 2*pad+1, n)
        oy = xp.asarray(oy) if _GPU else oy
        ox = xp.asarray(ox) if _GPU else ox
        ar = xp.arange(side)
        im = P[xp.arange(n)[:, None, None, None],
               xp.arange(3)[None, :, None, None],
               (oy[:, None]+ar[None, :])[:, None, :, None],
               (ox[:, None]+ar[None, :])[:, None, None, :]]
    fl = rg.random(n) < 0.5
    fl = xp.asarray(fl) if _GPU else fl
    im = xp.where(fl[:, None, None, None], im[:, :, :, ::-1], im)
    return im.reshape(n, -1)


class Model:
    """A folded convolution of c_out channels, its head, and its state.

    Carries the Adam moments so training can be paused and resumed, which
    is what lets a parent be joined mid-run and continued."""

    def __init__(self, g, c_in, k, c_out, nc, seed, v=None, H=None, b0=None,
                 bo=None):
        self.g, self.c_in, self.k, self.c_out = g, c_in, k, c_out
        idx, K, hid, macs = windowed(g, c_in, k, c_out)
        self.idx = to_dev(idx, np.int32) if _GPU else idx
        self.K, self.hid, self.macs = K, hid, macs
        self.D = c_in*g*g
        rg = np.random.default_rng(seed)
        if v is None:
            v = rg.normal(0, np.sqrt(2.0/(c_in*k*k)), K).astype(np.float32)
            v[-1] = 0.0
        if H is None:
            H = rg.normal(0, np.sqrt(2.0/hid), (hid, nc)).astype(np.float32)
        self.P = [to_dev(v), to_dev(H),
                  to_dev(np.zeros(hid) if b0 is None else b0),
                  to_dev(np.zeros(nc) if bo is None else bo)]
        self.M = [xp.zeros_like(p) for p in self.P]
        self.V = [xp.zeros_like(p) for p in self.P]
        self.t = 0

    def logits(self, x):
        W = self.P[0][self.idx].reshape(self.D, self.hid)
        h = xp.maximum(x @ W + self.P[2], 0)
        return h @ self.P[1] + self.P[3]

    def acc(self, Xte, yte):
        out = []
        for s in range(0, Xte.shape[0], 4096):
            out.append(to_host(self.logits(Xte[s:s+4096])))
        lg = np.concatenate(out)
        return float((lg.argmax(1) == yte).mean()), lg

    def step(self, Xtr, Ytr, cfg, epochs, side, rg):
        n = Xtr.shape[0]
        for ep in range(epochs):
            perm = rg.permutation(n)
            for st in range(0, n, cfg["batch"]):
                b = perm[st:st+cfg["batch"]]
                x = Xtr[b]
                if cfg["augment"]:
                    x = augment(x, side, cfg["aug_pad"], rg)
                y = Ytr[b]
                W = self.P[0][self.idx].reshape(self.D, self.hid)
                z = x @ W + self.P[2]; h = xp.maximum(z, 0)
                lg = h @ self.P[1] + self.P[3]
                e = xp.exp(lg - lg.max(1, keepdims=True))
                d = (e/e.sum(1, keepdims=True) - y)/len(b)
                d0 = (d @ self.P[1].T)*(z > 0)
                G = [scatter(x.T @ d0, self.idx, self.K), h.T @ d,
                     d0.sum(0), d.sum(0)]
                self.t += 1
                for i, (p_, gr) in enumerate(zip(self.P, G)):
                    self.M[i] = 0.9*self.M[i] + 0.1*gr
                    self.V[i] = 0.999*self.V[i] + 0.001*gr*gr
                    self.P[i] = p_ - cfg["lr"]*(self.M[i]/(1-0.9**self.t)) \
                        / (xp.sqrt(self.V[i]/(1-0.999**self.t))+1e-8)


def concatenate(parents, nc, seed):
    """A child holding every parent's channels, whose head is the parents'
    heads stacked and scaled.

    The scaling makes the child START at exactly its parents' average
    prediction, so anything it gains afterwards is co-adaptation rather
    than the ensemble it was built from."""
    p0 = parents[0]
    c_out = sum(p.c_out for p in parents)
    child = Model(p0.g, p0.c_in, p0.k, c_out, nc, seed)
    k, ci = p0.k, p0.c_in
    v = np.zeros(child.K, np.float32)
    off = 0
    Hs, b0s = [], []
    for p in parents:
        pv = to_host(p.P[0])
        for cin in range(ci):
            for co in range(p.c_out):
                src = (cin*p.c_out + co)*k*k
                dst = (cin*c_out + off + co)*k*k
                v[dst:dst+k*k] = pv[src:src+k*k]
        Hs.append(to_host(p.P[1])); b0s.append(to_host(p.P[2]))
        off += p.c_out
    n = len(parents)
    child.P[0] = to_dev(v)
    child.P[1] = to_dev(np.concatenate(Hs, 0)/n)
    child.P[2] = to_dev(np.concatenate(b0s, 0))
    child.P[3] = to_dev(np.mean([to_host(p.P[3]) for p in parents], 0))
    return child


CIFAR_MEAN = np.array([0.4914, 0.4822, 0.4465])
CIFAR_STD = np.array([0.2470, 0.2435, 0.2616])


def _verified_cifar(verbose=True):
    import glob, pickle
    for root in sorted(glob.glob("/kaggle/input/*")) + \
            ["/kaggle/working", "./cifar", "./data", "."]:
        if not os.path.isdir(root):
            continue
        cands = []
        for p in glob.glob(os.path.join(root, "**", "*"), recursive=True):
            if not os.path.isfile(p):
                continue
            try:
                if p.endswith(".npy"):
                    cands.append(np.load(p, allow_pickle=True))
                elif (p.endswith((".pickle", ".pkl", ".p"))
                      or os.path.basename(p).startswith(("data_batch",
                                                         "test_batch"))):
                    with open(p, "rb") as fh:
                        d = pickle.load(fh, encoding="bytes")
                    if isinstance(d, dict):
                        for kk, val in d.items():
                            kk = kk.decode() if isinstance(kk, bytes) else kk
                            a = np.asarray(val)
                            if a.size > 100 and kk in ("data", "labels",
                                                       "fine_labels", "x", "y"):
                                cands.append(a)
            except Exception:
                continue
        imgs = [a for a in cands if a.ndim >= 2 and len(a) >= 1000
                and a.size // len(a) == 3072]
        labs = [np.asarray(a).ravel() for a in cands
                if a.ndim <= 2 and np.issubdtype(np.asarray(a).dtype,
                                                 np.integer)
                and 1000 <= a.size <= 100000]
        if not imgs or not labs:
            continue
        X = np.concatenate([a.reshape(len(a), -1) for a in imgs])
        y = np.concatenate(labs)
        if len(y) != len(X):
            continue
        Xf = X.astype(np.float64)
        if Xf.max() > 1.5:
            Xf = Xf/255.0
        best = None
        for lay, shp in (("HWC", (-1, 32, 32, 3)), ("CHW", (-1, 3, 32, 32))):
            im = Xf[:2000].reshape(shp)
            mu = im.mean((0, 1, 2)) if lay == "HWC" else im.mean((0, 2, 3))
            sd = im.std((0, 1, 2)) if lay == "HWC" else im.std((0, 2, 3))
            e = float(np.abs(mu-CIFAR_MEAN).max() + np.abs(sd-CIFAR_STD).max())
            if best is None or e < best[0]:
                best = (e, lay)
        if best[0] > 0.05:
            continue
        if verbose:
            print(f"  verified CIFAR-10 at {root} ({best[1]}, {len(X):,} "
                  f"images)", flush=True)
        Xr = Xf.astype(np.float32)
        Xr = (Xr.reshape(-1, 32, 32, 3) if best[1] == "HWC"
              else Xr.reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1))
        return Xr, y.astype(np.int64)
    if verbose:
        print("  no verified CIFAR-10 found; downloading", flush=True)
    return None


def load(cfg):
    got = _verified_cifar()
    if got is not None:
        X, y = got
    else:
        from tensorflow import keras
        (a, b), (c, d) = keras.datasets.cifar10.load_data()
        X = np.concatenate([a, c]).astype(np.float32)/255.0
        y = np.concatenate([b, d]).ravel().astype(np.int64)
    side = 32
    if cfg["grid"] != side:
        s = side // cfg["grid"]
        X = X.reshape(-1, cfg["grid"], s, cfg["grid"], s, 3).mean(axis=(2, 4))
        side = cfg["grid"]
    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 = X[tr].mean((0, 1, 2), keepdims=True)
    sd = X[tr].std((0, 1, 2), keepdims=True) + 1e-8
    f = lambda Z: ((Z-mu)/sd).transpose(0, 3, 1, 2).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], side


CFG = dict(grid=16, c_in=3, k=3, base_channels=8, n_train=50000,
           batch=128, lr=1e-3, augment=True, aug_pad=2,
           seeds=(0, 1), join_at=30, total_epochs=90,
           # accumulate: the pool keeps every member, so widths grow
           #   8, 8 -> 16 -> 32
           # replace:  the child replaces its parents, so width is flat
           pool_mode="accumulate", rounds=2)


def main(**over):
    CFG.update(over)
    t0 = time.time()
    print("=" * 78)
    print("GROWING A MODEL FROM ITS OWN PARENTS, MID-TRAINING")
    print("=" * 78)
    print(f"  backend: {'cupy (GPU)' if _GPU else 'numpy (CPU)'}")
    for k_, v in CFG.items():
        print(f"  {k_:14s} = {v}")
    g, ci, k = CFG["grid"], CFG["c_in"], CFG["k"]
    C0 = CFG["base_channels"]
    print(f"\n  the child starts at EXACTLY its parents' ensemble, so what it")
    print(f"  gains afterwards is co-adaptation and not the join itself")
    print("=" * 78, flush=True)

    Xtr, Ytr, Xte, yte, side = load(CFG)
    Xtr, Ytr, Xte = to_dev(Xtr), to_dev(Ytr), to_dev(Xte)
    rg = np.random.default_rng(1234)
    res = {}

    # the parents, to the join point
    pool = [Model(g, ci, k, C0, 10, s) for s in CFG["seeds"]]
    for m in pool:
        m.step(Xtr, Ytr, CFG, CFG["join_at"], side, rg)
    for i, m in enumerate(pool):
        a, _ = m.acc(Xte, yte)
        print(f"  parent {i} at join ({CFG['join_at']} epochs): {a:.4f}"
              f"   [{time.time()-t0:.0f}s]", flush=True)

    # the child, and the ensemble it starts as
    child = concatenate(pool, 10, 99)
    ca, _ = child.acc(Xte, yte)
    lgs = [m.acc(Xte, yte)[1] for m in pool]
    ens = float((sum(lgs).argmax(1) == yte).mean())
    print(f"\n  child at birth {ca:.4f}, parents' ensemble {ens:.4f}"
          f"   (should agree to rounding)")

    rest = CFG["total_epochs"] - CFG["join_at"]
    child.step(Xtr, Ytr, CFG, rest, side, rg)
    for m in pool:
        m.step(Xtr, Ytr, CFG, rest, side, rg)
    ch, _ = child.acc(Xte, yte)
    pa = [m.acc(Xte, yte) for m in pool]
    ens_end = float((sum(p[1] for p in pa).argmax(1) == yte).mean())
    res["grown"] = dict(acc=ch, channels=child.c_out, values=child.K,
                        macs=child.macs)
    res["ensemble"] = dict(acc=ens_end, channels=sum(m.c_out for m in pool),
                           values=sum(m.K for m in pool),
                           macs=sum(m.macs for m in pool))
    for i, (a, _) in enumerate(pa):
        res[f"parent {i}"] = dict(acc=a, channels=pool[i].c_out,
                                  values=pool[i].K, macs=pool[i].macs)

    # THE control: the same width from scratch, same total epochs
    wide = Model(g, ci, k, child.c_out, 10, 777)
    wide.step(Xtr, Ytr, CFG, CFG["total_epochs"], side, rg)
    wa, _ = wide.acc(Xte, yte)
    res["wide from scratch"] = dict(acc=wa, channels=wide.c_out,
                                    values=wide.K, macs=wide.macs)
    print(f"  wide from scratch {wa:.4f}   [{time.time()-t0:.0f}s]",
          flush=True)

    # and the join-point sweep: late joining, same total budget
    late = max(CFG["join_at"] + 1,
               CFG["total_epochs"] - max(5, rest//4))
    pool2 = [Model(g, ci, k, C0, 10, s+50) for s in CFG["seeds"]]
    for m in pool2:
        m.step(Xtr, Ytr, CFG, late, side, rg)
    child2 = concatenate(pool2, 10, 98)
    child2.step(Xtr, Ytr, CFG, CFG["total_epochs"]-late, side, rg)
    la, _ = child2.acc(Xte, yte)
    res[f"joined late ({late})"] = dict(acc=la, channels=child2.c_out,
                                        values=child2.K, macs=child2.macs)
    json.dump(res, open("grown.json", "w"), indent=2)

    print("\n" + "=" * 78)
    print("  READOUT")
    print("=" * 78)
    print(f"  {'arm':>22s} {'channels':>9s} {'values':>8s} {'multiplies':>11s}"
          f" {'accuracy':>9s}")
    for nm, r in res.items():
        print(f"  {nm:>22s} {r['channels']:9d} {r['values']:8,} "
              f"{r['macs']:11,} {r['acc']:9.4f}")
    gr, wd = res["grown"]["acc"], res["wide from scratch"]["acc"]
    en = res["ensemble"]["acc"]
    bp = max(res[f"parent {i}"]["acc"] for i in range(len(pool)))
    print(f"\n  grown against the wide control  {gr-wd:+.4f}"
          f"   <- the one that matters")
    print(f"  grown against its own ensemble  {gr-en:+.4f}")
    print(f"  grown against the best parent   {gr-bp:+.4f}")
    print(f"  early join against late join    "
          f"{gr-res[f'joined late ({late})']['acc']:+.4f}")
    print()
    if gr > wd + 0.005 and gr > en + 0.005:
        print("  GROWING WINS. The child beats both the same width trained")
        print("  from scratch and the ensemble it was born as, so the")
        print("  parents' separate early learning is worth something AND the")
        print("  halves co-adapt after joining. A tournament has a basis.")
    elif gr > en + 0.005:
        print("  CO-ADAPTATION IS REAL BUT NOT WORTH THE ROUTE. The child")
        print("  beats the ensemble it started as, and does not beat the same")
        print("  width trained from scratch — so joining helps, and simply")
        print("  starting wide helps more.")
    elif abs(gr - en) < 0.005:
        print("  THE JOIN CHANGES NOTHING. The child ends where its parents'")
        print("  ensemble ends, so the halves are not co-adapting: they are")
        print("  two models sharing a head.")
    else:
        print("  GROWING LOSES. Read the table — the child is behind what it")
        print("  was built from, which means joining disturbed parents that")
        print("  were doing better apart.")
    print(f"\n  total {time.time()-t0:.0f}s; wrote grown.json")


if __name__ == "__main__":
    main()