| """Distillation training v2 — more data, longer schedule, sharded teacher cache. |
| |
| Changes vs v1: |
| - full RAVDESS train set: actors 1-18, speech (60/actor) + song (44/actor) = up to 1872 utts |
| (v1 used a 900-file speech-only cap) |
| - 2000 AdamW steps (v1: 900), same losses |
| - teacher caches precomputed in parallel on both T4s (precompute_shard.py) |
| - best checkpoint still selected on held-out validation slice of the train pool; |
| final quality gate remains eval_surgical.py on actors 19-24 |
| """ |
|
|
| import glob |
| import json |
| import math |
| import os |
| import random |
| 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 |
|
|
| 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" |
| OUT_DIR = "/tmp/work/surgical_out" |
| BATCH = 8 |
| LR = 3e-4 |
| STEPS = 2000 |
| WARMUP = 100 |
| SEED = 42 |
| VAL_FRAC = 0.08 |
| FP16 = True |
| DEV = "cuda:0" |
|
|
|
|
| def load_cache(): |
| shards = sorted(glob.glob(os.path.join(CACHE_DIR, "shard_*.pt"))) |
| assert shards, "no teacher shards found" |
| cache = [] |
| for s in shards: |
| cache.extend(torch.load(s, map_location="cpu", weights_only=False)) |
| return cache |
|
|
|
|
| def make_batch(pool): |
| 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) |
| inten = torch.zeros(B, T, 1) |
| 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() |
| inten[j, :t] = it["inten"].float() |
| lengths[j] = t |
| return states.to(DEV), probs.to(DEV), inten.to(DEV), lengths.to(DEV) |
|
|
|
|
| def compute_loss(enc, batch): |
| states, tprobs, tinten, lengths = batch |
| out = enc.adapter(states, lengths) |
| G = SurgicalConfig().group_size |
| embs, flen = out["frame_embs"], out["frame_lengths"] |
| T = states.size(1) |
| T_out = 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) |
|
|
| tinten_g = F.interpolate(tinten.transpose(1, 2), size=T_out, mode="area").transpose(1, 2) |
| l_int = (((out["frame_intensity"] - tinten_g) ** 2) * om).sum() / om.sum().clamp(min=1) |
|
|
| loss = (1.0 * l_frame + 0.5 * l_mse + 0.5 * l_utt_cos |
| + 1.0 * l_utt_kl + 1.0 * l_frm + 0.5 * l_int) |
| 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(), "inten": l_int.item(), |
| } |
|
|
|
|
| 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_v2.log"), "a", buffering=1) |
|
|
| def P(*a): |
| print(*a, flush=True); print(*a, file=log) |
|
|
| P(f"=== distill v2 start {time.strftime('%H:%M:%S')} device={DEV} ===") |
|
|
| enc = SurgicalEmotionEncoder.from_pretrained(FUNASR_DIR, device=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) |
|
|
| cache = load_cache() |
| P(f"teacher cache: {len(cache)} utts from {len(glob.glob(CACHE_DIR + '/shard_*.pt'))} shards") |
| random.shuffle(cache) |
| n_val = max(8, int(len(cache) * VAL_FRAC)) |
| val, train = cache[:n_val], cache[n_val:] |
| P(f"train={len(train)} val={len(val)}") |
|
|
| 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)) |
| 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 % 50 == 0: |
| P(f"step {step:4d} loss {loss.item():.4f} " + |
| " ".join(f"{k}={v:.3f}" for k, v in parts.items())) |
| if step % 200 == 0 or step == STEPS - 1: |
| adapter.eval() |
| vlosses = [] |
| with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16, enabled=FP16): |
| for _ in range(4): |
| _, vp = compute_loss(enc, make_batch(val)) |
| vlosses.append(sum(vp.values())) |
| vloss = float(np.mean(vlosses)) |
| P(f" >> val step {step}: {vloss:.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() |
|
|
