70 lines
1.8 KiB
Python
70 lines
1.8 KiB
Python
|
import time
|
||
|
|
from collections import namedtuple, OrderedDict
|
||
|
|
from typing import List, Dict, Any
|
||
|
|
|
||
|
|
import tabulate
|
||
|
|
import torch
|
||
|
|
import torch.distributed as dist
|
||
|
|
|
||
|
|
DeviceTime = namedtuple("DeviceTime", ["cpu", "gpu"])
|
||
|
|
|
||
|
|
|
||
|
|
class Functor:
|
||
|
|
|
||
|
|
def __init__(self, fn, *args, **kwargs):
|
||
|
|
self.fn = fn
|
||
|
|
self.args = args
|
||
|
|
self.kwargs = kwargs
|
||
|
|
|
||
|
|
def __call__(self):
|
||
|
|
return self.fn(*self.args, **self.kwargs)
|
||
|
|
|
||
|
|
|
||
|
|
def cuda_timeit(fn: Functor, dist_barrier=False) -> DeviceTime:
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
|
||
|
|
start = torch.cuda.Event(enable_timing=True)
|
||
|
|
stop = torch.cuda.Event(enable_timing=True)
|
||
|
|
start.record(torch.cuda.current_stream())
|
||
|
|
|
||
|
|
t0 = time.time()
|
||
|
|
fn()
|
||
|
|
t1 = time.time()
|
||
|
|
|
||
|
|
stop.record(torch.cuda.current_stream())
|
||
|
|
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
if dist_barrier:
|
||
|
|
dist.barrier()
|
||
|
|
|
||
|
|
gpu_time = start.elapsed_time(stop)
|
||
|
|
cpu_time = t1 - t0
|
||
|
|
return DeviceTime(cpu_time, gpu_time)
|
||
|
|
|
||
|
|
|
||
|
|
def cuda_benchmark(fn: Functor, num_repeated=10, num_warmup=1, dist_barrier=False) -> DeviceTime:
|
||
|
|
[fn() for _ in range(num_warmup)]
|
||
|
|
times = [cuda_timeit(fn, dist_barrier=dist_barrier) for _ in range(num_repeated)]
|
||
|
|
times.sort(key=lambda t: t.gpu)
|
||
|
|
|
||
|
|
if num_repeated >= 10:
|
||
|
|
times = times[3:-3]
|
||
|
|
|
||
|
|
avg_gpu_time = sum([t.gpu for t in times]) / len(times)
|
||
|
|
avg_cpu_time = sum([t.cpu for t in times]) / len(times)
|
||
|
|
|
||
|
|
return DeviceTime(avg_cpu_time * 1000, avg_gpu_time)
|
||
|
|
|
||
|
|
|
||
|
|
def show_benchmark_results(times: List[DeviceTime], extra_info: Dict[Any, List]=None):
|
||
|
|
data = extra_info or OrderedDict()
|
||
|
|
|
||
|
|
if len(times) != 0:
|
||
|
|
cpu_times = [round(t.cpu, 6) for t in times]
|
||
|
|
gpu_times = [round(t.gpu, 6) for t in times]
|
||
|
|
|
||
|
|
data["CPU Time(ms)"] = cpu_times
|
||
|
|
data["GPU Time(ms)"] = gpu_times
|
||
|
|
|
||
|
|
print(tabulate.tabulate(extra_info, headers=data.keys()))
|