Quick benchmark.
(fn, name, num_runs=5)
| 16 | |
| 17 | |
| 18 | def benchmark(fn, name, num_runs=5): |
| 19 | """Quick benchmark.""" |
| 20 | times = [] |
| 21 | for i in range(num_runs): |
| 22 | if torch.cuda.is_available(): |
| 23 | torch.cuda.synchronize() |
| 24 | start = time.perf_counter() |
| 25 | |
| 26 | result = fn() |
| 27 | |
| 28 | if torch.cuda.is_available(): |
| 29 | torch.cuda.synchronize() |
| 30 | elapsed = (time.perf_counter() - start) * 1000 |
| 31 | times.append(elapsed) |
| 32 | print(f" {name} run {i+1}: {elapsed:.1f}ms") |
| 33 | |
| 34 | return { |
| 35 | "name": name, |
| 36 | "mean": statistics.mean(times), |
| 37 | "std": statistics.stdev(times) if len(times) > 1 else 0, |
| 38 | } |
| 39 | |
| 40 | |
| 41 | def main(): |