| """Compare student v1 (900-utt/900-step distill) vs v2 (full-data/2000-step). |
| |
| Both checkpoints are loaded on top of the same frozen teacher; evaluated on the |
| identical held-out set (actors 19-24 speech, up to 4 files per emotion = 32 files). |
| Metrics per student vs teacher: top-1 agreement, JS divergence, pooled cosine, |
| frame cosine; plus intensity std (dynamics) and mean inference latency. |
| Saves /tmp/work/surgical_out/v1_v2_compare.json |
| """ |
|
|
| import glob |
| import json |
| import os |
| import sys |
| import time |
| import warnings |
|
|
| warnings.filterwarnings("ignore") |
| import logging |
|
|
| logging.disable(logging.CRITICAL) |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| import torchaudio |
|
|
| sys.path.insert(0, "/tmp/work/s2st_emotion") |
| from surgical_emotion import SurgicalConfig, SurgicalEmotionEncoder |
|
|
| FUNASR_DIR = "/tmp/work/emotion2vec_plus_base" |
| V1_CKPT = "/tmp/work/v1_student/surgical_emotion_module.pt" |
| V2_CKPT = "/tmp/work/surgical_out/surgical_emotion_module.pt" |
| OUT = "/tmp/work/surgical_out/v1_v2_compare.json" |
| DEV = "cuda:0" |
| LABELS = ["angry", "disgusted", "fearful", "happy", "neutral", "other", "sad", "surprised", "unk"] |
| EMO_FIELD = {"01": "neutral", "02": "calm", "03": "happy", "04": "sad", |
| "05": "angry", "06": "fearful", "07": "disgust", "08": "surprised"} |
|
|
| torch.set_grad_enabled(False) |
|
|
|
|
| def load_wav(path, max_s=12.0): |
| wav, sr = torchaudio.load(path) |
| wav = wav.mean(0) |
| if sr != 16000: |
| wav = torchaudio.functional.resample(wav, sr, 16000) |
| wav = wav[: int(max_s * 16000)] |
| return wav / wav.abs().max().clamp(min=1e-4) |
|
|
|
|
| def js_div(p, q, eps=1e-8): |
| p = p.clamp(min=eps); q = q.clamp(min=eps) |
| m = 0.5 * (p + q) |
| kl = lambda a, b: float((a * (a.log() - b.log())).sum(-1)) |
| return 0.5 * kl(p, m) + 0.5 * kl(q, m) |
|
|
|
|
| def rms_env(wav, T): |
| frames = torch.zeros(T) |
| for i in range(T): |
| seg = wav[i * 320: (i + 1) * 320] |
| frames[i] = float(seg.pow(2).mean().sqrt()) if seg.numel() else (frames[i - 1] if i else 0.0) |
| if T >= 9: |
| frames = F.conv1d(frames[None, None], torch.ones(9)[None, None] / 9.0, padding=4)[0, 0] |
| lo, hi = frames.min(), frames.max() |
| return torch.full_like(frames, 0.5) if hi - lo < 1e-6 else (frames - lo) / (hi - lo) |
|
|
|
|
| def main(): |
| fun = AutoModel if False else None |
| from funasr import AutoModel |
| fun = AutoModel(model=FUNASR_DIR, trust_remote_code=False, device=DEV, disable_update=True) |
| teacher = fun.model.eval() |
|
|
| v1 = SurgicalEmotionEncoder(teacher, SurgicalConfig()).to(DEV) |
| v1.adapter.load_state_dict(torch.load(V1_CKPT, map_location=DEV, weights_only=False)["adapter_state"]) |
| v1.eval() |
| v2 = SurgicalEmotionEncoder(teacher, SurgicalConfig()).to(DEV) |
| v2.adapter.load_state_dict(torch.load(V2_CKPT, map_location=DEV, weights_only=False)["adapter_state"]) |
| v2.eval() |
| students = {"v1": v1, "v2": v2} |
|
|
| files = sorted(glob.glob("/tmp/work/ravdess_kag/**/*.wav", recursive=True)) |
| files = [f for f in files if os.path.basename(f).split("-")[1] == "01" |
| and 19 <= int(os.path.basename(f).split("-")[-1].split(".")[0]) <= 24] |
| by_emo = {} |
| for f in files: |
| by_emo.setdefault(os.path.basename(f).split("-")[2], []).append(f) |
| bench = [] |
| for e in sorted(by_emo): |
| bench.extend(by_emo[e][:4]) |
| print(f"comparison inputs: {len(bench)} held-out files", flush=True) |
|
|
| report = {"n_files": len(bench), "per_student": {}, "files": {}} |
| stats = {k: {"agree": [], "js": [], "pool_cos": [], "frame_cos": [], "ms": [], |
| "inten_std": [], "inten_corr": []} for k in students} |
|
|
| for f in bench: |
| wav = load_wav(f).cpu() |
| dur = wav.numel() / 16000 |
| states, lengths = None, None |
| outs = {} |
| for name, enc in students.items(): |
| t0 = time.perf_counter() |
| outs[name] = enc(wav[None].to(DEV), torch.tensor([wav.numel()], device=DEV)) |
| torch.cuda.synchronize() |
| stats[name]["ms"].append((time.perf_counter() - t0) * 1000) |
| if states is None: |
| src = F.layer_norm(wav[None], wav[None].shape).to(DEV) |
| r = teacher.extract_features(src, padding_mask=None, mask=False) |
| states, lengths = r["x"].detach(), torch.tensor([r["x"].size(1)], device=DEV) |
| tp = torch.softmax(teacher.proj(states.mean(1)), -1)[0] |
|
|
| T = int(lengths[0]) |
| tgt = rms_env(wav, T) |
| tmax = torch.softmax(teacher.proj(states[0]), -1)[:, :8].max(-1).values.cpu().clamp(0, 1).sqrt() |
| tgt = tmax * tgt |
|
|
| row = {"dur_s": round(dur, 2), "ravdess": EMO_FIELD.get(os.path.basename(f).split("-")[2], "?"), |
| "teacher_top": LABELS[int(tp.argmax())]} |
| for name in students: |
| o = outs[name] |
| sp = o["utt_probs"][0] |
| a = int(sp.argmax().item() == tp.argmax().item()) |
| j = js_div(sp[None], tp[None]) |
| pc = float(F.cosine_similarity(o["utt_emb"][0], states.mean(1)[0], dim=0)) |
| up = F.interpolate(o["frame_embs"].transpose(1, 2), size=states.size(1), |
| mode="nearest").transpose(1, 2) |
| fc = float(F.cosine_similarity(up[0], states[0], dim=-1).mean()) |
| si = o["frame_intensity"][0, :, 0].cpu() |
| tgt_ds = F.interpolate(tgt[None, None], size=si.size(0), mode="linear")[0, 0] |
| ic = float(np.corrcoef(si.numpy(), tgt_ds.numpy())[0, 1]) if si.var() > 1e-6 else None |
| stats[name]["agree"].append(a); stats[name]["js"].append(j) |
| stats[name]["pool_cos"].append(pc); stats[name]["frame_cos"].append(fc) |
| stats[name]["inten_std"].append(float(si.std())) |
| if ic is not None: |
| stats[name]["inten_corr"].append(ic) |
| row[name] = {"top": LABELS[int(sp.argmax())], "p": round(float(sp.max()), 3), |
| "js": round(j, 5), "frame_cos": round(fc, 4)} |
| report["files"][os.path.basename(f)] = row |
|
|
| for name in students: |
| s = stats[name] |
| report["per_student"][name] = { |
| "top1_agreement": float(np.mean(s["agree"])), |
| "mean_js": float(np.mean(s["js"])), |
| "mean_pooled_cos": float(np.mean(s["pool_cos"])), |
| "mean_frame_cos": float(np.mean(s["frame_cos"])), |
| "mean_intensity_std": float(np.mean(s["inten_std"])), |
| "mean_intensity_corr": float(np.mean(s["inten_corr"])) if s["inten_corr"] else None, |
| "mean_ms_per_utt": float(np.mean(s["ms"])), |
| } |
| os.makedirs(os.path.dirname(OUT), exist_ok=True) |
| with open(OUT, "w") as fp: |
| json.dump(report, fp, indent=1) |
| print(json.dumps(report["per_student"], indent=1), flush=True) |
| print("COMPARE_DONE", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
