#!/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()