(model, batch_sizes, dataset)
| 362 | # |
| 363 | |
| 364 | def measure_execution_time(model, batch_sizes, dataset): |
| 365 | dataset_for_model = dataset.remove_columns(["example_id", "offset_mapping"]) |
| 366 | dataset_for_model.set_format("torch") |
| 367 | batch_size_to_time_sec = {} |
| 368 | for batch_size in batch_sizes: |
| 369 | batch = { |
| 370 | k: dataset_for_model[k][:batch_size].cuda() |
| 371 | for k in dataset_for_model.column_names |
| 372 | } |
| 373 | |
| 374 | with torch.no_grad(): |
| 375 | baseline_predictions = model(**batch) |
| 376 | timer = benchmark.Timer( |
| 377 | stmt="model(**batch)", globals={"model": model, "batch": batch} |
| 378 | ) |
| 379 | p50 = timer.blocked_autorange().median * 1000 |
| 380 | batch_size_to_time_sec[batch_size] = p50 |
| 381 | |
| 382 | model_c = torch.compile(model, fullgraph=True) |
| 383 | timer = benchmark.Timer( |
| 384 | stmt="model(**batch)", globals={"model": model_c, "batch": batch} |
| 385 | ) |
| 386 | p50 = timer.blocked_autorange().median * 1000 |
| 387 | batch_size_to_time_sec[f"{batch_size}_compile"] = p50 |
| 388 | new_predictions = model_c(**batch) |
| 389 | |
| 390 | return batch_size_to_time_sec |
| 391 | |
| 392 | |
| 393 |
no test coverage detected