File size: 3,755 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 | """Precompute frozen-teacher caches for a shard of training files.
Usage: python3 precompute_shard.py <device> <shard_idx> <n_shards>
Files are enumerated deterministically (train actors 1-18, speech + song);
this process handles every n_shards-th file starting at shard_idx.
Saves /tmp/work/teacher_cache/shard_<idx>.pt
"""
import glob
import os
import sys
import warnings
warnings.filterwarnings("ignore")
import logging
logging.disable(logging.CRITICAL)
import torch
import torchaudio
sys.path.insert(0, "/tmp/work/s2st_emotion")
from funasr import AutoModel
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 2
MAX_SEC = 12.0
OUT_DIR = "/tmp/work/teacher_cache"
os.makedirs(OUT_DIR, exist_ok=True)
def gather_files():
files = sorted(glob.glob("/tmp/work/ravdess_kag/**/*.wav", recursive=True))
out = []
for f in files:
b = os.path.basename(f)
parts = b.split("-")
try:
actor = int(parts[-1].split(".")[0])
channel = parts[1] # 01=speech, 02=song
except Exception:
continue
if actor <= 18 and channel in ("01", "02"):
out.append(f)
return out
def rms_envelope_50hz(wav, T_frames):
frames = torch.zeros(T_frames)
for i in range(T_frames):
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_frames >= 9:
frames = torch.nn.functional.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():
files = gather_files()
mine = files[SHARD::N_SHARDS]
print(f"[shard {SHARD}/{N_SHARDS} on {DEV}] total files={len(files)}, mine={len(mine)}", flush=True)
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)
cache, skipped = [], 0
for i, f in enumerate(mine):
try:
wav, sr = torchaudio.load(f)
wav = wav.mean(0)
if sr != 16000:
wav = torchaudio.functional.resample(wav, sr, 16000)
except Exception:
skipped += 1
continue
if wav.numel() < 8000:
skipped += 1
continue
if wav.numel() > MAX_SEC * 16000:
wav = wav[: int(MAX_SEC * 16000)]
wav = wav / wav.abs().max().clamp(min=1e-4)
states, lengths = enc.extract_teacher_states(wav[None].to(DEV),
torch.tensor([wav.numel()], device=DEV))
probs, inten = enc.teacher_frame_probs(states, lengths)
T = int(lengths[0].item())
rms = rms_envelope_50hz(wav, T)
inten = (inten.squeeze(0)[:, 0].cpu().clamp(min=0, max=1).sqrt() * rms).unsqueeze(-1)
cache.append({
"states": states.squeeze(0).half().cpu(),
"probs": probs.squeeze(0).half().cpu(),
"inten": inten.half().cpu(),
"file": os.path.basename(f),
})
if (i + 1) % 200 == 0:
print(f"[shard {SHARD}] {i+1}/{len(mine)}", flush=True)
torch.save(cache, f"{OUT_DIR}/shard_{SHARD}.pt")
print(f"[shard {SHARD}] saved {len(cache)} utts (skipped {skipped}) -> shard_{SHARD}.pt", flush=True)
if __name__ == "__main__":
main()
|