train_heads.py

#!/usr/bin/env python3
"""Train one small MLP regression head per dimension on frozen VoiceClap-commercial
embeddings.  Huber loss, bucket-stratified 10% val split, bucket-balanced sampling,
grid sweep over arch/lr/wd/dropout/delta, select per dimension by val Pearson r.

Head arch is byte-compatible with laion/voicenet-dimension-predictors-commercial
`MLPHead` for the 1-layer variants; 2-layer variants add an explicit arch tag.
The output affine (y = out*ysd + ymu) is FOLDED into the last Linear at export,
so a head consumes standardised embeddings and emits the score directly.

Env: DIM_SHARD / N_SHARD to split dimensions across array tasks.
"""
import os, json, glob, math, time, sys
import numpy as np, torch, torch.nn as nn

AD = "/e/data1/datasets/playground/mmlaion/schuhmann1/dramabox/attrdistill"
EMB = f"{AD}/embeddings"; HEADS = f"{AD}/heads"
os.makedirs(HEADS, exist_ok=True)
DEV = "cuda" if torch.cuda.is_available() else "cpu"
SHARD = int(os.environ.get("DIM_SHARD", "0")); NSHARD = int(os.environ.get("N_SHARD", "1"))
STEPS = int(os.environ.get("STEPS", "4000")); BS = int(os.environ.get("BS", "4096"))
EVAL_EVERY = int(os.environ.get("EVAL_EVERY", "250")); PATIENCE = int(os.environ.get("PATIENCE", "8"))
SEED = 1234

sys.path.insert(0, f"{AD}/scripts")
from embed_worker import SCORE_KEYS


# ----------------------------------------------------------------- data
def load_all():
    """Consolidated, key-de-duplicated containers written by consolidate.py."""
    E = np.load(f"{EMB}/emb_all.f16.npy")
    L = np.load(f"{EMB}/labels_all.f32.npy").copy()
    names = json.load(open(f"{EMB}/MANIFEST.json"))["label_columns"]
    # the embedder only ever sees the first 30 s, so longer durations are unknowable
    di = names.index("duration"); L[:, di] = np.minimum(L[:, di], 30.0)
    return E, L, names, None


def make_buckets(y, nb=10):
    """Equal-width bins over [p0.5,p99.5]; merge bins with <1% of data into neighbours."""
    lo, hi = np.percentile(y, 0.5), np.percentile(y, 99.5)
    if hi - lo < 1e-6: return None
    edges = np.linspace(lo, hi, nb + 1)
    b = np.clip(np.digitize(y, edges[1:-1]), 0, nb - 1)
    # merge tiny bins into the nearest larger neighbour
    for _ in range(nb):
        cnt = np.bincount(b, minlength=nb)
        alive = np.where(cnt > 0)[0]
        if len(alive) <= 2: break
        small = [i for i in alive if cnt[i] < 0.01 * len(y)]
        if not small: break
        i = min(small, key=lambda k: cnt[k])
        others = [k for k in alive if k != i]
        j = min(others, key=lambda k: abs(k - i))
        b[b == i] = j
    u, b = np.unique(b, return_inverse=True)
    return b


# ----------------------------------------------------------------- model
class Head1(nn.Module):
    """Identical topology to voicenet MLPHead: Linear-GELU-Dropout-Linear."""
    kind = "mlp1"
    def __init__(self, D, H, p):
        super().__init__()
        self.f1 = nn.Linear(D, H); self.act = nn.GELU(); self.dp = nn.Dropout(p); self.f2 = nn.Linear(H, 1)
    def forward(self, x): return self.f2(self.dp(self.act(self.f1(x))))
    def last(self): return self.f2

class Head2(nn.Module):
    kind = "mlp2ln"
    def __init__(self, D, H, H2, p):
        super().__init__()
        self.f1 = nn.Linear(D, H); self.ln = nn.LayerNorm(H); self.act = nn.GELU()
        self.dp = nn.Dropout(p); self.f2 = nn.Linear(H, H2); self.act2 = nn.GELU()
        self.dp2 = nn.Dropout(p); self.f3 = nn.Linear(H2, 1)
    def forward(self, x):
        h = self.dp(self.act(self.ln(self.f1(x))))
        return self.f3(self.dp2(self.act2(self.f2(h))))
    def last(self): return self.f3

ARCHS = [
    ("mlp1_h64",      lambda p: Head1(768, 64, p)),
    ("mlp1_h128",     lambda p: Head1(768, 128, p)),
    ("mlp2_h96_64",   lambda p: Head2(768, 96, 64, p)),
    ("mlp2_h64_32",   lambda p: Head2(768, 64, 32, p)),
]
def nparams(m): return sum(q.numel() for q in m.parameters())


def spearman(a, b):
    ra = np.argsort(np.argsort(a)).astype(np.float64)
    rb = np.argsort(np.argsort(b)).astype(np.float64)
    return float(np.corrcoef(ra, rb)[0, 1])

def pearson(a, b):
    if a.std() < 1e-9 or b.std() < 1e-9: return 0.0
    return float(np.corrcoef(a, b)[0, 1])


def run_one(Xtr, ytr, wtr, Xva, yva, arch_fn, lr, wd, drop, delta, ysd, ymu, gen):
    net = arch_fn(drop).to(DEV)
    opt = torch.optim.AdamW(net.parameters(), lr=lr, weight_decay=wd)
    sched = torch.optim.lr_scheduler.LambdaLR(
        opt, lambda s: min(1.0, (s + 1) / 200) * 0.5 * (1 + math.cos(math.pi * min(1.0, s / STEPS))))
    lossf = nn.HuberLoss(delta=delta)
    N = len(ytr)
    best = (-1e9, None, 0)
    bad = 0
    for step in range(STEPS):
        idx = torch.multinomial(wtr, BS, replacement=True, generator=gen)
        xb = Xtr[idx].float(); yb = ytr[idx]
        loss = lossf(net(xb).squeeze(-1), yb)
        opt.zero_grad(set_to_none=True); loss.backward(); opt.step(); sched.step()
        if (step + 1) % EVAL_EVERY == 0 or step == STEPS - 1:
            net.eval()
            with torch.no_grad():
                pv = torch.cat([net(Xva[i:i + 65536].float()).squeeze(-1) for i in range(0, len(Xva), 65536)])
            net.train()
            r = pearson(pv.cpu().numpy(), yva.cpu().numpy())
            if r > best[0] + 1e-4:
                best = (r, {k: v.detach().clone() for k, v in net.state_dict().items()}, step + 1); bad = 0
            else:
                bad += 1
                if bad >= PATIENCE: break
    net.load_state_dict(best[1]); net.eval()
    return net, best[0], best[2]


def evaluate(net, X, y, ysd, ymu, bkt=None):
    with torch.no_grad():
        p = torch.cat([net(X[i:i + 65536].float()).squeeze(-1) for i in range(0, len(X), 65536)]).cpu().numpy()
    yt = y.cpu().numpy()
    P = p * ysd + ymu; Y = yt * ysd + ymu
    out = {"mae": float(np.abs(P - Y).mean()), "rmse": float(np.sqrt(((P - Y) ** 2).mean())),
           "r": pearson(P, Y), "rho": spearman(P, Y),
           "pred_std": float(P.std()), "true_std": float(Y.std())}
    if bkt is not None:                              # bucket-balanced val metrics
        rng = np.random.default_rng(0); sel = []
        cnt = np.bincount(bkt); m = int(max(20, np.median(cnt[cnt > 0])))
        for b in np.unique(bkt):
            ix = np.where(bkt == b)[0]
            sel.append(rng.choice(ix, size=min(len(ix), m), replace=False))
        sel = np.concatenate(sel)
        out["r_bal"] = pearson(P[sel], Y[sel]); out["mae_bal"] = float(np.abs(P[sel] - Y[sel]).mean())
    return out, P, Y


def main():
    t0 = time.time()
    E, L, names, _ = load_all()
    print(f"loaded {E.shape[0]} samples, {E.shape[1]}-d, {L.shape[1]} label columns", flush=True)
    rng = np.random.default_rng(SEED)
    Xall = torch.from_numpy(E).to(DEV)                       # float16 (N,768)
    # feature standardisation on a large random subsample (same for every head)
    sub = rng.choice(len(E), size=min(200_000, len(E)), replace=False)
    mu = Xall[torch.from_numpy(sub).to(DEV)].float().mean(0)
    sd = Xall[torch.from_numpy(sub).to(DEV)].float().std(0).clamp_min(1e-6)
    json.dump({"n": int(len(E))}, open(f"{HEADS}/_n.json", "w"))

    mydims = [i for i in range(len(names)) if i % NSHARD == SHARD]
    results = {}
    for di in mydims:
        name = names[di]
        y = L[:, di]
        ok = np.isfinite(y)
        if ok.sum() < 2000:
            print(f"SKIP {name}: only {int(ok.sum())} labelled", flush=True)
            results[name] = {"skipped": f"only {int(ok.sum())} labelled samples"}; continue
        yi = y[ok]; idx_all = np.where(ok)[0]
        bkt = make_buckets(yi)
        if bkt is None:
            print(f"SKIP {name}: degenerate range", flush=True)
            results[name] = {"skipped": "degenerate score range"}; continue
        # stratified 10% val per bucket
        va = np.zeros(len(yi), bool)
        for b in np.unique(bkt):
            ix = np.where(bkt == b)[0]
            va[rng.choice(ix, size=max(1, int(round(0.1 * len(ix)))), replace=False)] = True
        tr = ~va
        ymu, ysd = float(yi[tr].mean()), float(yi[tr].std() + 1e-8)
        yn = (yi - ymu) / ysd
        gi = torch.from_numpy(idx_all).to(DEV)
        Xtr = Xall[gi[torch.from_numpy(np.where(tr)[0]).to(DEV)]]
        Xva = Xall[gi[torch.from_numpy(np.where(va)[0]).to(DEV)]]
        Xtr = ((Xtr.float() - mu) / sd).half(); Xva = ((Xva.float() - mu) / sd).half()
        ytr = torch.from_numpy(yn[tr]).float().to(DEV); yva = torch.from_numpy(yn[va]).float().to(DEV)
        cnt = np.bincount(bkt[tr]); w = 1.0 / np.maximum(cnt[bkt[tr]], 1)
        wtr = torch.from_numpy(w).float().to(DEV)
        gen = torch.Generator(device=DEV); gen.manual_seed(SEED + di)

        grid = [(a, lr, wd, 0.1, 1.0) for a in ARCHS for lr in (1e-3, 3e-3) for wd in (1e-4, 1e-2)]
        best = None
        for (aname, afn), lr, wd, drop, delta in grid:
            net, r, st = run_one(Xtr, ytr, wtr, Xva, yva, afn, lr, wd, drop, delta, ysd, ymu, gen)
            if best is None or r > best[1]:
                best = (net, r, dict(arch=aname, lr=lr, wd=wd, drop=drop, delta=delta, steps=st))
        # refine the winning arch
        bafn = dict(ARCHS)[best[2]["arch"]]
        for drop, delta in ((0.0, 1.0), (0.1, 0.5), (0.2, 1.0)):
            net, r, st = run_one(Xtr, ytr, wtr, Xva, yva, bafn, best[2]["lr"], best[2]["wd"], drop, delta, ysd, ymu, gen)
            if r > best[1]:
                best = (net, r, dict(best[2], drop=drop, delta=delta, steps=st))
        net, _, cfg = best
        m, P, Y = evaluate(net, Xva, yva, ysd, ymu, bkt=bkt[va])
        npar = nparams(net)
        # fold output affine into the last Linear so the head emits the raw score
        with torch.no_grad():
            lastl = net.last(); lastl.weight.mul_(ysd); lastl.bias.mul_(ysd).add_(ymu)
        res = dict(name=name, n_train=int(tr.sum()), n_val=int(va.sum()), n_buckets=int(bkt.max() + 1),
                   arch=cfg["arch"], params=npar, lr=cfg["lr"], wd=cfg["wd"], dropout=cfg["drop"],
                   huber_delta=cfg["delta"], steps=cfg["steps"], y_mean=ymu, y_std=ysd,
                   y_min=float(yi.min()), y_max=float(yi.max()), **m)
        results[name] = res
        torch.save({"state_dict": net.state_dict(), "kind": net.kind, "arch": cfg["arch"],
                    "meta": res}, f"{HEADS}/{name.replace('/','_')}.pt")
        # small scatter sample for the report
        k = min(3000, len(P)); s = rng.choice(len(P), k, replace=False)
        np.save(f"{HEADS}/_scatter_{name.replace('/','_')}.npy",
                np.stack([P[s], Y[s]]).astype(np.float16))
        print(f"[{name}] {cfg['arch']} {npar/1e3:.1f}k  ntr={res['n_train']} nva={res['n_val']}  "
              f"MAE={m['mae']:.3f} RMSE={m['rmse']:.3f} r={m['r']:.3f} rho={m['rho']:.3f} "
              f"r_bal={m.get('r_bal',0):.3f}  [{time.time()-t0:.0f}s]", flush=True)
        json.dump(results, open(f"{HEADS}/_results_shard{SHARD}.json", "w"), indent=1)
    np.save(f"{HEADS}/_norm.npy", np.stack([mu.cpu().numpy(), sd.cpu().numpy()]))
    json.dump(results, open(f"{HEADS}/_results_shard{SHARD}.json", "w"), indent=1)
    print(f"SHARD {SHARD} DONE in {time.time()-t0:.0f}s", flush=True)


if __name__ == "__main__":
    main()