| |
| from __future__ import annotations |
| import time |
| from statistics import median |
| import numpy as np |
| import torch |
| import ssl_mamba |
|
|
| _CUT = 64600 |
|
|
| def time_batch_sizes(): |
| print("device", torch.cuda.get_device_name()) |
| m = ssl_mamba.SSLMamba() |
| m.load() |
| cuda = getattr(m, "device", "cpu") == "cuda" |
| results = {} |
| print("=== batch-size sweep ===") |
| for bs in [1, 2, 4, 8, 16, 24, 32]: |
| audios = [np.random.randn(_CUT).astype(np.float32) for _ in range(bs)] |
| srs = [16000] * bs |
| try: |
| m.score_batch(audios, srs) |
| if cuda: torch.cuda.synchronize() |
| reps = [] |
| for _ in range(5): |
| if cuda: torch.cuda.synchronize() |
| t0 = time.perf_counter() |
| m.score_batch(audios, srs) |
| if cuda: torch.cuda.synchronize() |
| reps.append(time.perf_counter() - t0) |
| med = median(reps); ups = bs / med |
| results[bs] = ups |
| print(f"bs={bs:3d} median={med*1000:8.3f} ms utt/s={ups:8.2f}") |
| except torch.cuda.OutOfMemoryError: |
| print(f"bs={bs:3d} OOM (skipped)") |
| if cuda: torch.cuda.empty_cache() |
| m.unload() |
| best = max(results, key=results.get) |
| print(f"best batch_size = {best} (utt/s={results[best]:.2f})") |
| return best, results |
|
|
| if __name__ == "__main__": |
| time_batch_sizes() |
|
|
