File size: 6,436 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 163 164 165 166 167 168 169 | """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()
|