wehub-resource-sync
48 lines
1.2 KiB
Python
48 lines
1.2 KiB
Python
from __future__ import annotations
|
|
|
|
import time
|
|
import typing as t
|
|
|
|
import numpy as np
|
|
from rich.console import Console
|
|
from rich.table import Table
|
|
|
|
P = t.ParamSpec("P")
|
|
R = t.TypeVar("R")
|
|
OrigFunc = t.Callable[P, R]
|
|
DecoratedFunc = t.Callable[P, tuple[np.floating, np.floating]]
|
|
|
|
|
|
def timeit(func: OrigFunc, iteration: int = 3) -> DecoratedFunc:
|
|
def function_timer(
|
|
*args: P.args, **kwargs: P.kwargs
|
|
) -> tuple[np.floating, np.floating]:
|
|
"""
|
|
Time the execution of a function and returns the time taken
|
|
"""
|
|
# warmup
|
|
func(*args, **kwargs)
|
|
|
|
runtimes = []
|
|
for _ in range(iteration):
|
|
start = time.time()
|
|
# we dont care about the return value
|
|
func(*args, **kwargs)
|
|
end = time.time()
|
|
runtime = end - start
|
|
runtimes.append(runtime)
|
|
|
|
return np.mean(runtimes), np.var(runtimes)
|
|
|
|
return function_timer
|
|
|
|
|
|
def print_table(result):
|
|
table = Table("Batch Name", "(mean, var)", title="Benchmark Results")
|
|
|
|
for batch_name, (mean, var) in result.items():
|
|
table.add_row(batch_name, f"{mean:.4f}, {var:.4f}")
|
|
|
|
console = Console()
|
|
console.print(table)
|