| """ |
| 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, |
| |
| |
| |
| 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 = {} |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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() |
|
|
