Benchmark model inference. Returns (output, median_latency_ms). Uses CUDA events for precise GPU timing.
(
model: nn.Module,
model_input: Union[torch.Tensor, Dict[str, torch.Tensor]],
warmup: int = WARMUP_RUNS,
timed: int = TIMED_RUNS,
)
| 256 | # --------------------------------------------------------------------------- |
| 257 | |
| 258 | def benchmark_model( |
| 259 | model: nn.Module, |
| 260 | model_input: Union[torch.Tensor, Dict[str, torch.Tensor]], |
| 261 | warmup: int = WARMUP_RUNS, |
| 262 | timed: int = TIMED_RUNS, |
| 263 | ) -> Tuple[Any, float]: |
| 264 | """ |
| 265 | Benchmark model inference. Returns (output, median_latency_ms). |
| 266 | Uses CUDA events for precise GPU timing. |
| 267 | """ |
| 268 | if not torch.cuda.is_available(): |
| 269 | raise RuntimeError("CUDA is required for benchmarking.") |
| 270 | |
| 271 | def _run(): |
| 272 | with torch.no_grad(): |
| 273 | if isinstance(model_input, dict): |
| 274 | return model(**model_input) |
| 275 | else: |
| 276 | return model(model_input) |
| 277 | |
| 278 | # Warmup |
| 279 | print(f" Warmup: {warmup} runs...", end="", flush=True) |
| 280 | for _ in range(warmup): |
| 281 | output = _run() |
| 282 | torch.cuda.synchronize() |
| 283 | print(" done") |
| 284 | |
| 285 | # Timed runs |
| 286 | print(f" Timed: {timed} runs...", end="", flush=True) |
| 287 | start_events = [torch.cuda.Event(enable_timing=True) for _ in range(timed)] |
| 288 | end_events = [torch.cuda.Event(enable_timing=True) for _ in range(timed)] |
| 289 | |
| 290 | torch.cuda.synchronize() |
| 291 | for i in range(timed): |
| 292 | start_events[i].record() |
| 293 | _run() |
| 294 | end_events[i].record() |
| 295 | torch.cuda.synchronize() |
| 296 | print(" done") |
| 297 | |
| 298 | # Compute median |
| 299 | times_ms = sorted(s.elapsed_time(e) for s, e in zip(start_events, end_events)) |
| 300 | median_ms = times_ms[len(times_ms) // 2] |
| 301 | |
| 302 | # Final reference output (deterministic) |
| 303 | with torch.no_grad(): |
| 304 | output = _run() |
| 305 | torch.cuda.synchronize() |
| 306 | |
| 307 | return output, median_ms |
| 308 | |
| 309 | |
| 310 | # --------------------------------------------------------------------------- |