Files
livekit-cameras/scripts/spike_bench_cpu_metricgan.py
root 42443420be chore: OpenVINO MetricGAN GPU spike (do not merge)
Arc B580 lost the kill test: batch-14 mask p50 103 ms vs Torch CPU
70 ms enhance / OpenVINO CPU 5 ms. Production stays on known-good
CPU workers (tag known-good-cpu-2026-08-28). Hop stays 200 ms.
2026-08-28 07:18:54 +00:00

77 lines
2.3 KiB
Python

#!/usr/bin/env python3
"""CPU baseline: thread-capped MetricGAN enhance_batch (project .venv)."""
from __future__ import annotations
import argparse
import os
import statistics
import time
os.environ.setdefault("OMP_NUM_THREADS", "1")
os.environ.setdefault("MKL_NUM_THREADS", "1")
os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
import torch
from speechbrain.inference.enhancement import SpectralMaskEnhancement
SAVEDIR = "/root/.cache/speechbrain-enhancement"
def pct(xs, p):
xs = sorted(xs)
i = int(round((p / 100.0) * (len(xs) - 1)))
return xs[max(0, min(i, len(xs) - 1))]
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--batch", type=int, default=14)
ap.add_argument("--samples", type=int, default=16000)
ap.add_argument("--iters", type=int, default=30)
ap.add_argument("--warmup", type=int, default=5)
args = ap.parse_args()
torch.set_num_threads(1)
torch.set_num_interop_threads(1)
m = SpectralMaskEnhancement.from_hparams(
source="speechbrain/metricgan-plus-voicebank",
savedir=SAVEDIR,
run_opts={"device": "cpu"},
)
wav = torch.randn(args.batch, args.samples)
lengths = torch.ones(args.batch)
with torch.no_grad():
for _ in range(args.warmup):
m.enhance_batch(wav, lengths=lengths)
times = []
for _ in range(args.iters):
t0 = time.perf_counter()
m.enhance_batch(wav, lengths=lengths)
times.append((time.perf_counter() - t0) * 1000.0)
feats = m.compute_features(wav)
net = m.mods.enhance_model
for _ in range(args.warmup):
net(feats, lengths=lengths)
mask_times = []
for _ in range(args.iters):
t0 = time.perf_counter()
net(feats, lengths=lengths)
mask_times.append((time.perf_counter() - t0) * 1000.0)
def report(name, xs):
per = [x / args.batch for x in xs]
print(
f"{name} batch={args.batch} n={len(xs)} "
f"p50={statistics.median(xs):.2f}ms p95={pct(xs,95):.2f}ms "
f"p99={pct(xs,99):.2f}ms "
f"per_cam_p50={statistics.median(per):.2f}ms "
f"per_cam_p95={pct(per,95):.2f}ms"
)
report("cpu_enhance_batch", times)
report("cpu_mask_only", mask_times)
if __name__ == "__main__":
main()