| """Throughput of the local build against the padded-bank conv1d formulation. |
| |
| python dev/bench.py |
| """ |
| import os |
| import sys |
| import time |
|
|
| import torch |
| import torch.nn.functional as Fn |
|
|
| ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| sys.path.insert(0, ROOT) |
| import load_local |
|
|
| rp = load_local.load() |
|
|
| PAIRS = [("44.1k -> 16k (speech in)", 44100, 16000), |
| ("48k -> 16k (speech in)", 48000, 16000), |
| ("16k -> 24k (tts out)", 16000, 24000), |
| ("44.1k -> 48k (device)", 44100, 48000)] |
| SECONDS = 30.0 |
| REPS = 25 |
|
|
|
|
| def timeit(fn, reps=REPS): |
| for _ in range(5): |
| fn() |
| best = 1e30 |
| for _ in range(reps): |
| t0 = time.perf_counter() |
| fn() |
| best = min(best, time.perf_counter() - t0) |
| return best |
|
|
|
|
| def ref(x, plan): |
| L, M, width, kern = plan.L, plan.M, plan.width, plan.kernel |
| B, T = x.shape |
| pad = Fn.pad(x, (width, width + M)) |
| r = Fn.conv1d(pad[:, None], kern[:, None, :], stride=M) |
| return r.transpose(1, 2).reshape(B, -1)[:, :plan.out_len(T)] |
|
|
|
|
| print(f"threads={torch.get_num_threads()} audio={SECONDS:.0f}s mono\n") |
| print(f"{'conversion':28} {'resample-poly':>13} {'conv1d bank':>12} " |
| f"{'speedup':>8} {'taps/out':>9} {'path':>6}") |
|
|
| for name, orig, new in PAIRS: |
| T = int(orig * SECONDS) |
| g = torch.Generator().manual_seed(0) |
| x = torch.randn(1, T, generator=g, dtype=torch.float32) |
| plan = rp.Resampler(orig, new) |
|
|
| t_mine = timeit(lambda: plan(x)) |
| t_ref = timeit(lambda: ref(x, plan)) |
| a, b = plan(x), ref(x, plan) |
| rel = ((a - b).abs().max() / b.abs().max()).item() |
| assert rel < 2e-5, rel |
| print(f"{name:28} {t_mine*1e3:12.2f}ms {t_ref*1e3:11.2f}ms " |
| f"{t_ref/t_mine:7.2f}x {plan.taps_per_output:9.1f} " |
| f"{'fused' if plan.fused else 'bank':>6}") |
|
|
| plan = rp.Resampler(44100, 16000) |
| x = torch.randn(1, int(44100 * SECONDS), dtype=torch.float32) |
| t = timeit(lambda: plan(x)) |
| print(f"\n44.1k->16k realtime factor: {SECONDS / t:,.0f}x " |
| f"({t*1e3:.2f} ms for {SECONDS:.0f} s)") |
| print(f"taps stored {plan.taps.numel()} of {plan.kernel.numel()} " |
| f"padded-bank entries ({plan.density*100:.1f}%)") |
|
|
