File size: 4,355 Bytes
6fdd8c7 | 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 | """Precompute v4 caches — same GPU-heavy pipeline as precompute_fast.py, but the
cache stores the RAW (unnormalized) 50 Hz RMS envelope instead of the old
min-max "conf-gated" intensity target. The v4 intensity target is derived in
train_v4 from (probs, rms_raw) with a global reference level.
Cache item: {states (T,768) f16, probs (T,9) f16, rms (T,) f16, file}
Usage: python3 precompute_v4.py <device> <shard_idx> <n_shards>
"""
import os, sys, threading, queue, warnings
warnings.filterwarnings("ignore")
import logging; logging.disable(logging.CRITICAL)
import torch, torchaudio
import torch.nn.functional as F
sys.path.insert(0, "/tmp/work/s2st_emotion")
from precompute_shard_v3 import gather_files, MAX_SEC
from surgical_emotion import SurgicalConfig, SurgicalEmotionEncoder
DEV = sys.argv[1] if len(sys.argv) > 1 else "cuda:0"
SHARD = int(sys.argv[2]) if len(sys.argv) > 2 else 0
N_SHARDS = int(sys.argv[3]) if len(sys.argv) > 3 else 4
N_LOAD = 4
OUT_DIR = "/tmp/work/teacher_cache_v4"
os.makedirs(OUT_DIR, exist_ok=True)
resamplers = {}
def to_16k(wav, sr):
if sr == 16000:
return wav
if sr not in resamplers:
resamplers[sr] = torchaudio.transforms.Resample(sr, 16000).to(DEV)
return resamplers[sr](wav.to(DEV)).cpu()
def load_one(f):
try:
wav, sr = torchaudio.load(f)
wav = wav.mean(0)
if sr != 16000:
wav = to_16k(wav, sr)
except Exception:
return None
if wav.numel() < 8000:
return None
if wav.numel() > MAX_SEC * 16000:
wav = wav[: int(MAX_SEC * 16000)]
wav = wav / wav.abs().max().clamp(min=1e-4)
return wav
def rms_raw_gpu(wav_dev, T_f):
"""RAW 50 Hz RMS envelope (9-tap box smoothed), NO min-max normalization."""
L = wav_dev.numel()
n_full = (L + 319) // 320
w = torch.zeros(T_f * 320, device=wav_dev.device, dtype=wav_dev.dtype)
w[: min(L, T_f * 320)] = wav_dev[: min(L, T_f * 320)]
ms = w.view(T_f, 320).pow(2).mean(1).sqrt()
if T_f > n_full:
ms[n_full:] = ms[n_full - 1]
if T_f >= 9:
ms = F.conv1d(ms[None, None], torch.ones(9, device=DEV)[None, None] / 9.0,
padding=4)[0, 0]
return ms
def main():
files, counts = gather_files()
mine = files[SHARD::N_SHARDS]
print(f"[v4 shard {SHARD}/{N_SHARDS} on {DEV}] total={len(files)} mine={len(mine)}", flush=True)
partial_path = f"{OUT_DIR}/shard_{SHARD}_partial.pt"
cache = []
if os.path.exists(partial_path):
cache = torch.load(partial_path, map_location="cpu", weights_only=False)
print(f"[v4 shard {SHARD}] resuming with {len(cache)} done", flush=True)
mine = mine[len(cache):]
from funasr import AutoModel
fun = AutoModel(model="/tmp/work/emotion2vec_plus_base", trust_remote_code=False,
device=DEV, disable_update=True)
enc = SurgicalEmotionEncoder(fun.model, SurgicalConfig()).to(DEV)
enc.eval()
torch.set_grad_enabled(False)
q = queue.Queue(maxsize=32)
def loader(paths, q):
for p in paths:
q.put((p, load_one(p)))
q.put(None)
th = threading.Thread(target=loader, args=(mine, q), daemon=True)
th.start()
skipped = 0
done = len(cache)
while True:
item = q.get()
if item is None:
break
path, wav = item
if wav is None:
skipped += 1
continue
wg = wav.to(DEV)
src = F.layer_norm(wg[None], wg[None].shape)
st = enc.teacher.extract_features(src, padding_mask=None, mask=False)["x"][0].float()
Ti = st.size(0)
probs = torch.softmax(enc.teacher.proj(st), dim=-1)
rms = rms_raw_gpu(wg, Ti)
cache.append({
"states": st.half().cpu(),
"probs": probs.half().cpu(),
"rms": rms.half().cpu(),
"file": os.path.basename(path),
})
done = len(cache)
if done % 500 == 0:
print(f"[v4 shard {SHARD}] {done}/{len(files)}", flush=True)
if done % 2000 == 0:
torch.save(cache, partial_path)
torch.save(cache, f"{OUT_DIR}/shard_{SHARD}.pt")
if os.path.exists(partial_path):
os.remove(partial_path)
print(f"[v4 shard {SHARD}] saved {len(cache)} utts (skipped {skipped})", flush=True)
if __name__ == "__main__":
main()
|