shawn
0e77adb4c3
- Full FFT overlap-add IR convolution in IRLoader (process(), set_mix(), toggle) - Lazy FFT computation — IR FFT padded to correct block+ir size on first process() - Wet/dry mix control, enabled/disabled toggle with tail clearing - Fixed pipeline._apply_ir_cab() to delegate to IRLoader.process() instead of poking internals (old code had array-size mismatch bug: IR FFT at ir_len vs block FFT at conv_size) - 46 tests: loading, convolution correctness, overlap-add state, mix, toggle, directory listing, performance budget (all <5ms even at 8192 taps), edge cases - scripts/download_irs.sh: free IR pack downloader (God's Cab, Seacow)
120 lines
4.0 KiB
Python
120 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
||
"""NAM inference benchmark — measure latency at 48kHz with various model sizes.
|
||
|
||
Usage:
|
||
python scripts/benchmark_nam.py # all models, all block sizes
|
||
python scripts/benchmark_nam.py --quick # default model only, 256-block
|
||
python scripts/benchmark_nam.py --model convnet --block 256
|
||
|
||
Results are extrapolated for RPi 4B (Cortex-A72) using a 4-8x slowdown factor
|
||
vs x86_64 PyTorch CPU.
|
||
"""
|
||
|
||
import argparse
|
||
import time
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
|
||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||
|
||
SAMPLE_RATE = 48000
|
||
RPI_SLOWDOWN_LOWER = 4.0 # Conservative RPi slowdown estimate
|
||
RPI_SLOWDOWN_UPPER = 8.0 # Pessimistic RPi slowdown estimate
|
||
|
||
|
||
def build_models():
|
||
"""Return list of (name, factory, params_estimate)."""
|
||
from nam.models.conv_net import ConvNet
|
||
|
||
models = [
|
||
("nano (4ch, 3 dil)", lambda: ConvNet(
|
||
channels=4, dilations=[1, 2, 4],
|
||
batchnorm=False, activation="Tanh",
|
||
)),
|
||
("feather (8ch, 5 dil)", lambda: ConvNet(
|
||
channels=8, dilations=[1, 2, 4, 8, 16],
|
||
batchnorm=False, activation="Tanh",
|
||
)),
|
||
("lite (12ch, 6 dil)", lambda: ConvNet(
|
||
channels=12, dilations=[1, 2, 4, 8, 16, 32],
|
||
batchnorm=False, activation="Tanh",
|
||
)),
|
||
("standard (16ch, 8 dil)", lambda: ConvNet(
|
||
channels=16, dilations=[1, 2, 4, 8, 16, 32, 64, 128],
|
||
batchnorm=False, activation="Tanh",
|
||
)),
|
||
("default (32ch, 4 dil)", lambda: ConvNet(
|
||
channels=32, dilations=[1, 2, 4, 8],
|
||
batchnorm=False, activation="Tanh",
|
||
)),
|
||
]
|
||
return models
|
||
|
||
|
||
def benchmark_model(model_fn, block_size, n_warmup=10, n_runs=500):
|
||
"""Run model on block_size samples, return average ms."""
|
||
import torch
|
||
|
||
model = model_fn()
|
||
model.eval()
|
||
with torch.no_grad():
|
||
x = torch.randn(1, block_size)
|
||
for _ in range(n_warmup):
|
||
_ = model(x)
|
||
|
||
start = time.perf_counter()
|
||
for _ in range(n_runs):
|
||
_ = model(x)
|
||
elapsed = time.perf_counter() - start
|
||
|
||
avg_ms = elapsed / n_runs * 1000.0
|
||
params = sum(p.numel() for p in model.parameters())
|
||
return avg_ms, params
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="NAM inference benchmark")
|
||
parser.add_argument("--quick", action="store_true", help="Only default model, 256-block")
|
||
parser.add_argument("--model", type=str, default=None, help="Model name substring to filter")
|
||
parser.add_argument("--block", type=int, default=None, help="Block size (override)")
|
||
args = parser.parse_args()
|
||
|
||
models = build_models()
|
||
block_sizes = [128, 256, 512]
|
||
|
||
if args.quick:
|
||
block_sizes = [256]
|
||
models = [m for m in models if "default" in m[0]]
|
||
|
||
if args.block:
|
||
block_sizes = [args.block]
|
||
|
||
if args.model:
|
||
models = [m for m in models if args.model.lower() in m[0].lower()]
|
||
|
||
print(f"{'Model':<30} {'Params':>8} {'Block':>6} {'x86 ms':>8} {'RPi ms (4x)':>12} {'RPi ms (8x)':>12} {'Budget%':>8}")
|
||
print("-" * 90)
|
||
|
||
for name, factory in models:
|
||
for bs in block_sizes:
|
||
block_ms = bs / SAMPLE_RATE * 1000.0
|
||
avg_ms, params = benchmark_model(factory, bs)
|
||
|
||
rpi_low = avg_ms * RPI_SLOWDOWN_LOWER
|
||
rpi_high = avg_ms * RPI_SLOWDOWN_UPPER
|
||
budget_pct = (avg_ms / block_ms) * 100
|
||
|
||
ok = "✓" if rpi_low < block_ms else "⚠" if rpi_high < block_ms else "✗"
|
||
|
||
print(f"{ok} {name:<27} {params:>8,} {bs:>6} {avg_ms:>7.3f}ms {rpi_low:>8.2f}ms {rpi_high:>8.2f}ms {budget_pct:>6.1f}%")
|
||
|
||
print()
|
||
print("Legend: ✓ = safe on RPi (4x est < block budget), ⚠ = marginal, ✗ = too slow")
|
||
print(f"Block budget: {block_sizes[0] / SAMPLE_RATE * 1000:.2f}ms @ {SAMPLE_RATE}Hz")
|
||
print(f"RPi estimate: {RPI_SLOWDOWN_LOWER:.0f}x-{RPI_SLOWDOWN_UPPER:.0f}x × x86 PyTorch CPU time")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |