| """Distillation training v5 — v4 recipe + weak supervision from REAL intensity labels. |
| |
| v4 hit the ceiling of its label-free proxy: the TARGET itself separates RAVDESS |
| normal(01)/strong(02) by only Cohen-d 0.23 (student 0.26). v5 adds utterance-level |
| pairwise ranking supervision from intensity labels that already exist in the |
| training corpora but were never used: |
| RAVDESS (speech+song) 03-0v-e-i-s-r-a.wav -> field i: 01=normal 02=strong |
| CREMA-D spk_utt_emo_lev.wav -> lev: LO=0, MD=0.5, HI=1 (XX skipped) |
| |
| Supervision = pairwise margin ranking on the utterance-mean intensity. It is |
| shift-invariant, so it cannot fight the frame-level distillation target; it only |
| enforces ORDERING (strong > normal, HI > MD > LO). Every 2nd batch is sampled |
| from the labeled subset so pairs are frequent. |
| |
| Same v4 cache (states/probs/rms — labels parsed from "file" at load), same |
| held-out gate (RAVDESS actors 19-24), G=8 (6.25 Hz), 4000 AdamW steps, batch 8. |
| """ |
| import glob, json, math, os, random, sys, time, warnings |
| warnings.filterwarnings("ignore") |
| import logging; logging.disable(logging.CRITICAL) |
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
|
|
| sys.path.insert(0, "/tmp/work/s2st_emotion") |
| from surgical_emotion import SurgicalConfig, SurgicalEmotionEncoder |
|
|
| FUNASR_DIR = "/tmp/work/emotion2vec_plus_base" |
| CACHE_DIR = "/tmp/work/teacher_cache_v4" |
| OUT_DIR = "/tmp/work/surgical_out_v5" |
| BATCH, LR, STEPS, WARMUP, SEED = 8, 3e-4, 4000, 200, 42 |
| VAL_FRAC, FP16, DEV = 0.05, True, "cuda:0" |
| EMO_IDX = [0, 1, 2, 3, 6, 7] |
| GATE_FLOOR = 0.02 |
| MARGIN = 0.10 |
| LAM_RANK = 1.0 |
| LABELED_PROB = 0.5 |
|
|
|
|
| def parse_label(fname): |
| """RAVDESS 01->0.0, 02->1.0 | CREMA-D LO->0.0, MD->0.5, HI->1.0 | else None.""" |
| stem = os.path.splitext(fname)[0] |
| if stem.count("-") >= 6: |
| p = stem.split("-") |
| if p[1] in ("01", "02"): |
| return {"01": 0.0, "02": 1.0}.get(p[3]) |
| if "_" in stem: |
| p = stem.split("_") |
| if len(p) == 4: |
| return {"LO": 0.0, "MD": 0.5, "HI": 1.0}.get(p[3]) |
| return None |
|
|
|
|
| def load_cache(): |
| shards = sorted(glob.glob(os.path.join(CACHE_DIR, "shard_*.pt"))) |
| assert shards, "no v4 teacher shards found" |
| cache = [] |
| for s in shards: |
| cache.extend(torch.load(s, map_location="cpu", weights_only=False)) |
| for it in cache: |
| it["label"] = parse_label(it.get("file", "")) |
| return cache |
|
|
|
|
| def make_batch(pool, labeled): |
| if labeled and random.random() < LABELED_PROB: |
| pool = labeled |
| items = random.sample(pool, min(BATCH, len(pool))) |
| T = max(it["states"].size(0) for it in items) |
| B = len(items) |
| states = torch.zeros(B, T, 768); probs = torch.zeros(B, T, 9); rms = torch.zeros(B, T) |
| labels = torch.zeros(B) |
| lengths = torch.zeros(B, dtype=torch.long) |
| for j, it in enumerate(items): |
| t = it["states"].size(0) |
| states[j, :t] = it["states"].float(); probs[j, :t] = it["probs"].float() |
| rms[j, :t] = it["rms"].float(); lengths[j] = t |
| labels[j] = it["label"] if it["label"] is not None else -1.0 |
| return states.to(DEV), probs.to(DEV), rms.to(DEV), lengths.to(DEV), labels.to(DEV) |
|
|
|
|
| def masked_pearson(a, b, m): |
| n = m.sum(1).clamp(min=1) |
| ma = (a * m).sum(1) / n; mb = (b * m).sum(1) / n |
| da = (a - ma[:, None]) * m; db = (b - mb[:, None]) * m |
| cov = (da * db).sum(1) |
| va = (da * da).sum(1); vb = (db * db).sum(1) |
| corr = cov / (va * vb).clamp(min=1e-10).sqrt() |
| return corr, va, vb, n |
|
|
|
|
| def intensity_target(tprobs, trms, global_ref): |
| p6 = tprobs[..., EMO_IDX].max(-1).values |
| env = (trms / global_ref).clamp(0.0, 1.0) |
| return torch.clamp(p6, min=GATE_FLOOR) * env |
|
|
|
|
| def compute_loss(enc, batch, global_ref): |
| states, tprobs, trms, lengths, labels = batch |
| out = enc.adapter(states, lengths) |
| embs, flen = out["frame_embs"], out["frame_lengths"] |
| T, T_out = states.size(1), embs.size(1) |
|
|
| up = F.interpolate(embs.transpose(1, 2), size=T, mode="nearest").transpose(1, 2) |
| m = (torch.arange(T, device=states.device)[None, :] < lengths[:, None]).unsqueeze(-1) |
| cos = F.cosine_similarity(up, states, dim=-1).unsqueeze(-1) |
| l_frame = ((1 - cos) * m).sum() / m.sum().clamp(min=1) |
| l_mse = F.mse_loss(F.layer_norm(up, (768,)), F.layer_norm(states, (768,)), reduction="none") |
| l_mse = (l_mse * m).sum() / m.sum().clamp(min=1) |
|
|
| om = (torch.arange(T_out, device=states.device)[None, :] < flen[:, None]).unsqueeze(-1) |
| utt_s = (embs * om).sum(1) / flen.unsqueeze(-1).float() |
| utt_probs_s = torch.softmax(enc.adapter.emotion_head(utt_s.float()), dim=-1) |
| utt_t = (states * m).sum(1) / lengths.unsqueeze(-1).float() |
| l_utt_cos = (1 - F.cosine_similarity(utt_s, utt_t, dim=-1)).mean() |
| utt_probs_t = torch.softmax(enc.teacher.proj(utt_t), dim=-1) |
| l_utt_kl = F.kl_div(torch.log(utt_probs_s.clamp(min=1e-8)), utt_probs_t, reduction="batchmean") |
|
|
| tprobs_g = F.interpolate(tprobs.transpose(1, 2), size=T_out, mode="area").transpose(1, 2) |
| tprobs_g = tprobs_g / tprobs_g.sum(-1, keepdim=True).clamp(min=1e-8) |
| l_frm = F.kl_div(torch.log(out["frame_probs"].clamp(min=1e-8)), tprobs_g, |
| reduction="none").sum(-1) |
| l_frm = (l_frm * om.squeeze(-1)).sum() / om.sum().clamp(min=1) |
|
|
| tgt50 = intensity_target(tprobs, trms, global_ref) |
| tgt = F.interpolate(tgt50[:, None, :], size=T_out, mode="area")[:, 0, :] |
| omf = om.squeeze(-1).float() |
| si = out["frame_intensity"][:, :, 0].float() |
| l_int_mse = (((si - tgt) ** 2) * omf).sum() / omf.sum().clamp(min=1) |
| corr, va, vb, n = masked_pearson(si, tgt, omf) |
| ok = (va > 1e-8) & (vb > 1e-8) |
| if ok.any(): |
| std_s = (va[ok] / n[ok]).clamp(min=1e-12).sqrt() |
| std_t = (vb[ok] / n[ok]).clamp(min=1e-12).sqrt() |
| l_corr = (1 - corr[ok]).mean() |
| l_var = ((std_s / std_t).clamp(1e-4, 1e4).log() ** 2).mean() |
| else: |
| l_corr = si.sum() * 0.0; l_var = si.sum() * 0.0 |
| l_int = l_int_mse + 1.0 * l_corr + 0.5 * l_var |
|
|
| |
| utt_pred = (si * omf).sum(1) / omf.sum(1).clamp(min=1) |
| has = labels >= 0 |
| n_pairs = 0 |
| l_rank = si.sum() * 0.0 |
| if has.sum() >= 2: |
| pl = utt_pred[has]; ql = labels[has] |
| diff = ql[:, None] - ql[None, :] |
| mask = diff > 0 |
| n_pairs = int(mask.sum()) |
| if n_pairs: |
| pr = pl[:, None] - pl[None, :] |
| l_rank = (F.relu(MARGIN - pr) * mask).sum() / mask.sum() |
|
|
| loss = (1.0 * l_frame + 0.5 * l_mse + 0.5 * l_utt_cos |
| + 1.0 * l_utt_kl + 1.0 * l_frm + 1.0 * l_int + LAM_RANK * l_rank) |
| return loss, {"frame": l_frame.item(), "mse": l_mse.item(), "utt_cos": l_utt_cos.item(), |
| "utt_kl": l_utt_kl.item(), "frm_kl": l_frm.item(), |
| "int_mse": l_int_mse.item(), |
| "int_corr": (1 - l_corr.item()) if ok.any() else float("nan"), |
| "int_var": l_var.item(), "rank": l_rank.item(), "npair": n_pairs} |
|
|
|
|
| def main(): |
| random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED) |
| os.makedirs(OUT_DIR, exist_ok=True) |
| log = open(os.path.join(OUT_DIR, "train_v5.log"), "a", buffering=1) |
| def P(*a): |
| print(*a, flush=True); print(*a, file=log) |
|
|
| P(f"=== distill v5 start {time.strftime('%H:%M:%S')} device={DEV} G=8 (6.25 Hz) + label ranking ===") |
|
|
| cache = load_cache() |
| labeled = [it for it in cache if it["label"] is not None] |
| P(f"teacher cache: {len(cache)} utts | labeled: {len(labeled)}") |
| nl = sum(1 for it in labeled if it["file"].count("-") >= 6) |
| P(f" ravdess-labeled={nl} cremad-labeled={len(labeled) - nl}") |
| random.shuffle(cache) |
| n_val = max(16, int(len(cache) * VAL_FRAC)) |
| val, train = cache[:n_val], cache[n_val:] |
| val_lab = [it for it in val if it["label"] is not None] |
| train_lab = [it for it in train if it["label"] is not None] |
| P(f"train={len(train)} (labeled {len(train_lab)}) val={len(val)} (labeled {len(val_lab)})") |
|
|
| p90s = [torch.quantile(it["rms"].float(), 0.90).item() for it in cache if it["rms"].numel()] |
| global_ref = float(np.median(p90s)) |
| P(f"GLOBAL_REF (median P90 rms): {global_ref:.6f}") |
| with open(os.path.join(OUT_DIR, "target_meta.json"), "w") as f: |
| json.dump({"global_ref": global_ref, "emo_idx": EMO_IDX, "gate_floor": GATE_FLOOR, |
| "formula": "tgt = clamp(max(p[emo6]),floor) * clip(rms/GLOBAL_REF,0,1)", |
| "supervision": f"pairwise margin ranking (margin={MARGIN}), labeled-batch prob={LABELED_PROB}", |
| "version": 5, "group_size": 8, "frame_rate_hz": 6.25}, f, indent=1) |
|
|
| enc = SurgicalEmotionEncoder.from_pretrained(FUNASR_DIR, device=DEV) |
| enc.cfg.group_size = 8 |
| enc.adapter = type(enc.adapter)(enc.cfg).to(DEV) |
| adapter = enc.adapter |
| opt = torch.optim.AdamW(adapter.parameters(), lr=LR, weight_decay=0.01) |
| sched = torch.optim.lr_scheduler.LambdaLR( |
| opt, lambda s: min((s + 1) / WARMUP, 0.5 * (1 + math.cos(math.pi * s / STEPS)))) |
| scaler = torch.amp.GradScaler("cuda", enabled=FP16) |
|
|
| best_val = float("inf") |
| for step in range(STEPS): |
| adapter.train() |
| with torch.autocast("cuda", dtype=torch.float16, enabled=FP16): |
| loss, parts = compute_loss(enc, make_batch(train, train_lab), global_ref) |
| opt.zero_grad(set_to_none=True) |
| scaler.scale(loss).backward() |
| scaler.unscale_(opt) |
| torch.nn.utils.clip_grad_norm_(adapter.parameters(), 1.0) |
| scaler.step(opt); scaler.update(); sched.step() |
| if step % 100 == 0: |
| P(f"step {step:4d} loss {loss.item():.4f} " + |
| " ".join(f"{k}={v:.3f}" for k, v in parts.items() if k != "npair") + |
| f" npair={parts['npair']}") |
| if step % 400 == 0 or step == STEPS - 1: |
| adapter.eval() |
| vlosses, vcorr, vrank = [], [], [] |
| with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16, enabled=FP16): |
| for _ in range(6): |
| _, vp = compute_loss(enc, make_batch(val, val_lab), global_ref) |
| vlosses.append(sum(v for k, v in vp.items() if k not in ("int_corr", "npair"))) |
| vcorr.append(vp["int_corr"]); vrank.append(vp["rank"]) |
| vloss = float(np.mean(vlosses)) |
| P(f" >> val step {step}: {vloss:.4f} val_int_corr={np.nanmean(vcorr):.3f} " |
| f"val_rank={np.mean(vrank):.4f}") |
| if vloss < best_val: |
| best_val = vloss |
| enc.save_surgical(os.path.join(OUT_DIR, "surgical_emotion_module.pt")) |
| P(" >> saved best") |
|
|
| enc.save_surgical(os.path.join(OUT_DIR, "surgical_emotion_module_last.pt")) |
| P(f"=== done, best_val={best_val:.4f} ===") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
