File size: 6,774 Bytes
1419b86 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 | """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()
|