| """Full teacher-vs-student benchmark on identical inputs. |
| |
| Paths measured per identical input: |
| teacher_pure : stock compute path (layer_norm -> extract_features -> mean -> proj -> softmax) |
| teacher_pipeline : stock funasr user path (AutoModel.generate, granularity=frame) |
| student_total : SurgicalEmotionEncoder end-to-end (frozen backbone + adapter, 12.5 Hz outputs) |
| adapter_only : adapter on precomputed teacher states (marginal surgical cost) |
| |
| Quality (same inputs, both models): |
| top-1 agreement, JS divergence of utterance probs, pooled-embedding cosine, |
| frame-level cosine (student upsampled to 50 Hz vs teacher states), |
| intensity stats + correlation with conf-gated energy target. |
| |
| Scaling: durations 1/2/3/4/10 s, adapter batch 1/4/8, fp16 autocast variant, |
| small CPU comparison. Saves benchmark_report.json. |
| """ |
|
|
| import glob |
| import json |
| import os |
| import sys |
| import time |
| import warnings |
|
|
| warnings.filterwarnings("ignore") |
| import logging |
|
|
| logging.disable(logging.CRITICAL) |
|
|
| sys.path.insert(0, "/tmp/work/s2st_emotion") |
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| import torchaudio |
|
|
| from funasr import AutoModel |
| from surgical_emotion import SurgicalConfig, SurgicalEmotionEncoder |
|
|
| FUNASR_DIR = "/tmp/work/emotion2vec_plus_base" |
| CKPT = "/tmp/work/surgical_out/surgical_emotion_module.pt" |
| OUT = "/tmp/work/surgical_out/benchmark_report.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"} |
|
|
| report = {"gpu": torch.cuda.get_device_name(0), "torch": torch.__version__, |
| "device": DEV, "labels": LABELS, "paths": {}, "files": {}, "scaling": {}, "quality_summary": {}} |
|
|
|
|
| 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 timeit(fn, warmup=3, reps=10): |
| for _ in range(warmup): |
| fn() |
| torch.cuda.synchronize() |
| ts = [] |
| for _ in range(reps): |
| t0 = time.perf_counter() |
| fn() |
| torch.cuda.synchronize() |
| ts.append(time.perf_counter() - t0) |
| return {"mean_ms": 1000 * float(np.mean(ts)), "std_ms": 1000 * float(np.std(ts)), |
| "min_ms": 1000 * float(np.min(ts))} |
|
|
|
|
| print("=== loading models (shared teacher backbone) ===", flush=True) |
| fun = AutoModel(model=FUNASR_DIR, trust_remote_code=False, device=DEV, disable_update=True) |
| teacher = fun.model.eval() |
| enc = SurgicalEmotionEncoder(teacher, SurgicalConfig()).to(DEV) |
| payload = torch.load(CKPT, map_location=DEV, weights_only=False) |
| enc.adapter.load_state_dict(payload["adapter_state"]) |
| enc.eval() |
| n_teacher = sum(p.numel() for p in teacher.parameters()) |
| n_adapter = sum(p.numel() for p in enc.adapter.parameters()) |
| report["params"] = {"teacher_backbone": n_teacher, "student_adapter": n_adapter, |
| "adapter_overhead_pct": 100.0 * n_adapter / n_teacher} |
| print(f"teacher params {n_teacher/1e6:.2f}M, adapter {n_adapter/1e6:.2f}M", flush=True) |
| torch.set_grad_enabled(False) |
|
|
| torch.cuda.reset_peak_memory_stats() |
|
|
| |
| files = sorted(glob.glob("/tmp/work/ravdess_kag/Actor_*/*.wav")) |
| 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_files = [] |
| for e in sorted(by_emo): |
| bench_files.extend(by_emo[e][:3]) |
| print(f"bench inputs: {len(bench_files)} files", flush=True) |
|
|
|
|
| |
| def teacher_pure(wav): |
| src = F.layer_norm(wav[None], wav[None].shape) |
| out = teacher.extract_features(src.to(DEV), padding_mask=None, mask=False) |
| x = out["x"] |
| probs = torch.softmax(teacher.proj(x.mean(1)), -1) |
| return x, probs |
|
|
|
|
| def teacher_pipeline(path): |
| return fun.generate(path, granularity="frame", extract_embedding=True) |
|
|
|
|
| def student_total(wav): |
| return enc(wav[None].to(DEV), torch.tensor([wav.numel()], device=DEV)) |
|
|
|
|
| def adapter_only(states, lengths): |
| return enc.from_states(states, lengths) |
|
|
|
|
| |
| print("=== per-file run (teacher & student on identical wavs) ===", flush=True) |
| agrees, jss, pool_coss, frame_coss, corrs = [], [], [], [], [] |
| per_path_times = {"teacher_pure": [], "student_total": [], "adapter_only": []} |
|
|
| for f in bench_files: |
| wav = load_wav(f).cpu() |
| dur = wav.numel() / 16000 |
|
|
| |
| t0 = time.perf_counter(); tx, tprob = teacher_pure(wav); torch.cuda.synchronize() |
| per_path_times["teacher_pure"].append(time.perf_counter() - t0) |
| states = tx.detach() |
| lengths = torch.tensor([states.size(1)], device=DEV) |
|
|
| t0 = time.perf_counter(); sout = student_total(wav); torch.cuda.synchronize() |
| per_path_times["student_total"].append(time.perf_counter() - t0) |
|
|
| t0 = time.perf_counter(); aout = adapter_only(states, lengths); torch.cuda.synchronize() |
| per_path_times["adapter_only"].append(time.perf_counter() - t0) |
|
|
| |
| sp = sout["utt_probs"][0] |
| tp = tprob[0] |
| agrees.append(int(sp.argmax().item() == tp.argmax().item())) |
| jss.append(js_div(sp[None], tp[None])) |
| pool_coss.append(float(F.cosine_similarity(sout["utt_emb"][0], states.mean(1)[0], dim=0))) |
| up = F.interpolate(sout["frame_embs"].transpose(1, 2), size=states.size(1), mode="nearest").transpose(1, 2) |
| frame_coss.append(float(F.cosine_similarity(up[0], states[0], dim=-1).mean())) |
|
|
| inten = sout["frame_intensity"][0, :, 0].cpu() |
| 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 |
| tgt_ds = F.interpolate(tgt[None, None], size=inten.size(0), mode="linear")[0, 0] |
| c = float(np.corrcoef(inten.numpy(), tgt_ds.numpy())[0, 1]) if inten.var() > 1e-6 and tgt_ds.var() > 1e-6 else None |
| if c is not None: |
| corrs.append(c) |
|
|
| report["files"][os.path.basename(f)] = { |
| "emotion": EMO_FIELD.get(os.path.basename(f).split("-")[2], "?"), |
| "dur_s": round(dur, 2), |
| "teacher_top": LABELS[int(tp.argmax())], "teacher_p": round(float(tp.max()), 3), |
| "student_top": LABELS[int(sp.argmax())], "student_p": round(float(sp.max()), 3), |
| "js": round(jss[-1], 5), |
| "pooled_cos": round(pool_coss[-1], 5), |
| "frame_cos": round(frame_coss[-1], 5), |
| "intensity": {"mean": round(float(inten.mean()), 3), "std": round(float(inten.std()), 3), |
| "min": round(float(inten.min()), 3), "max": round(float(inten.max()), 3), |
| "corr_target": round(c, 3) if c is not None else None}, |
| } |
|
|
| n = len(bench_files) |
| report["quality_summary"] = { |
| "n_files": n, |
| "top1_agreement": float(np.mean(agrees)), |
| "mean_js": float(np.mean(jss)), |
| "mean_pooled_cos": float(np.mean(pool_coss)), |
| "mean_frame_cos": float(np.mean(frame_coss)), |
| "mean_intensity_corr": float(np.mean(corrs)), |
| } |
| durs = np.array([load_wav(f).numel() / 16000 for f in bench_files]) |
| for k, v in per_path_times.items(): |
| v = np.array(v) |
| report["paths"][k] = { |
| "mean_ms_per_utt": float(1000 * v.mean()), "rtf_mean": float((v / durs).mean()), |
| "utt_per_s": float(1.0 / v.mean()), "audio_x_realtime": float((durs / v).mean()), |
| } |
| print(json.dumps(report["quality_summary"], indent=1), flush=True) |
|
|
| |
| print("=== duration scaling ===", flush=True) |
| base = load_wav(bench_files[0]) |
| long_wav = base.repeat(4)[: 10 * 16000] |
| report["scaling"]["duration"] = {} |
| for sec in [1, 2, 3, 4, 10]: |
| w = long_wav[: sec * 16000] |
| row = {"dur_s": sec} |
| row["teacher_pure"] = timeit(lambda: teacher_pure(w)) |
| row["student_total"] = timeit(lambda: student_total(w)) |
| st, ln = (lambda r: (r[0].detach(), torch.tensor([r[0].size(1)], device=DEV)))(teacher_pure(w)) |
| row["adapter_only"] = timeit(lambda: adapter_only(st, ln)) |
| for k in ["teacher_pure", "student_total", "adapter_only"]: |
| row[k]["rtf"] = row[k]["mean_ms"] / 1000.0 / sec |
| report["scaling"]["duration"][f"{sec}s"] = row |
| print(f"{sec}s: teacher {row['teacher_pure']['mean_ms']:.1f}ms (RTF {row['teacher_pure']['rtf']:.3f}) | " |
| f"student {row['student_total']['mean_ms']:.1f}ms (RTF {row['student_total']['rtf']:.3f}) | " |
| f"adapter {row['adapter_only']['mean_ms']:.2f}ms", flush=True) |
|
|
| |
| t0 = time.perf_counter() |
| for f in bench_files[:3]: |
| teacher_pipeline(f) |
| pipe_t = (time.perf_counter() - t0) / 3 |
| d3 = float(np.mean([load_wav(f).numel() / 16000 for f in bench_files[:3]])) |
| report["paths"]["teacher_pipeline_funasr"] = {"mean_ms_per_utt": pipe_t * 1000, "rtf_mean": pipe_t / d3, |
| "utt_per_s": 1.0 / pipe_t, "audio_x_realtime": d3 / pipe_t} |
|
|
| |
| print("=== adapter batch scaling ===", flush=True) |
| w3 = long_wav[: 3 * 16000] |
| x3, _ = teacher_pure(w3) |
| T3 = x3.size(1) |
| report["scaling"]["adapter_batch"] = {} |
| for B in [1, 4, 8]: |
| st = x3.expand(B, T3, 768).contiguous() |
| ln = torch.full((B,), T3, device=DEV, dtype=torch.long) |
| r = timeit(lambda: adapter_only(st, ln)) |
| r["utt_per_s"] = 1000.0 / r["mean_ms"] * B |
| report["scaling"]["adapter_batch"][f"B{B}"] = r |
| print(f"B={B}: {r['mean_ms']:.2f}ms -> {r['utt_per_s']:.1f} utt/s", flush=True) |
|
|
| |
| with torch.autocast("cuda", dtype=torch.float16): |
| r16 = timeit(lambda: student_total(w3)) |
| report["scaling"]["student_fp16_autocast_3s"] = r16 |
| print(f"student fp16 3s: {r16['mean_ms']:.1f}ms", flush=True) |
|
|
| |
| report["memory"] = {"peak_gpu_allocated_mb": torch.cuda.max_memory_allocated() / 1e6, |
| "peak_gpu_reserved_mb": torch.cuda.max_memory_reserved() / 1e6} |
|
|
| |
| print("=== CPU comparison (3 s input, 3 reps) ===", flush=True) |
| enc_cpu = SurgicalEmotionEncoder(teacher.to("cpu"), SurgicalConfig()) |
| enc_cpu.adapter.load_state_dict(payload["adapter_state"]) |
| enc_cpu.eval() |
| w3c = w3.cpu() |
|
|
|
|
| def student_cpu(): |
| with torch.no_grad(): |
| enc_cpu(w3c[None], torch.tensor([w3c.numel()])) |
|
|
|
|
| with torch.no_grad(): |
| tc = timeit(student_cpu, warmup=1, reps=3) |
| report["scaling"]["cpu_3s"] = {"student_total": tc} |
| print(f"cpu student 3s: {tc['mean_ms']:.0f}ms", flush=True) |
| teacher.to(DEV) |
| enc.to(DEV) |
|
|
| with open(OUT, "w") as fp: |
| json.dump(report, fp, indent=1) |
| print("BENCH_DONE", flush=True) |
|
|
