|
|
|
import logging |
|
import numpy as np |
|
from itertools import count |
|
from typing import List, Tuple |
|
import torch |
|
import tqdm |
|
from fvcore.common.timer import Timer |
|
|
|
from detectron2.utils import comm |
|
|
|
from .build import build_batch_data_loader |
|
from .common import DatasetFromList, MapDataset |
|
from .samplers import TrainingSampler |
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
class _EmptyMapDataset(torch.utils.data.Dataset): |
|
""" |
|
Map anything to emptiness. |
|
""" |
|
|
|
def __init__(self, dataset): |
|
self.ds = dataset |
|
|
|
def __len__(self): |
|
return len(self.ds) |
|
|
|
def __getitem__(self, idx): |
|
_ = self.ds[idx] |
|
return [0] |
|
|
|
|
|
def iter_benchmark( |
|
iterator, num_iter: int, warmup: int = 5, max_time_seconds: float = 60 |
|
) -> Tuple[float, List[float]]: |
|
""" |
|
Benchmark an iterator/iterable for `num_iter` iterations with an extra |
|
`warmup` iterations of warmup. |
|
End early if `max_time_seconds` time is spent on iterations. |
|
|
|
Returns: |
|
float: average time (seconds) per iteration |
|
list[float]: time spent on each iteration. Sometimes useful for further analysis. |
|
""" |
|
num_iter, warmup = int(num_iter), int(warmup) |
|
|
|
iterator = iter(iterator) |
|
for _ in range(warmup): |
|
next(iterator) |
|
timer = Timer() |
|
all_times = [] |
|
for curr_iter in tqdm.trange(num_iter): |
|
start = timer.seconds() |
|
if start > max_time_seconds: |
|
num_iter = curr_iter |
|
break |
|
next(iterator) |
|
all_times.append(timer.seconds() - start) |
|
avg = timer.seconds() / num_iter |
|
return avg, all_times |
|
|
|
|
|
class DataLoaderBenchmark: |
|
""" |
|
Some common benchmarks that help understand perf bottleneck of a standard dataloader |
|
made of dataset, mapper and sampler. |
|
""" |
|
|
|
def __init__( |
|
self, |
|
dataset, |
|
*, |
|
mapper, |
|
sampler=None, |
|
total_batch_size, |
|
num_workers=0, |
|
max_time_seconds: int = 90, |
|
): |
|
""" |
|
Args: |
|
max_time_seconds (int): maximum time to spent for each benchmark |
|
other args: same as in `build.py:build_detection_train_loader` |
|
""" |
|
if isinstance(dataset, list): |
|
dataset = DatasetFromList(dataset, copy=False, serialize=True) |
|
if sampler is None: |
|
sampler = TrainingSampler(len(dataset)) |
|
|
|
self.dataset = dataset |
|
self.mapper = mapper |
|
self.sampler = sampler |
|
self.total_batch_size = total_batch_size |
|
self.num_workers = num_workers |
|
self.per_gpu_batch_size = self.total_batch_size // comm.get_world_size() |
|
|
|
self.max_time_seconds = max_time_seconds |
|
|
|
def _benchmark(self, iterator, num_iter, warmup, msg=None): |
|
avg, all_times = iter_benchmark(iterator, num_iter, warmup, self.max_time_seconds) |
|
if msg is not None: |
|
self._log_time(msg, avg, all_times) |
|
return avg, all_times |
|
|
|
def _log_time(self, msg, avg, all_times, distributed=False): |
|
percentiles = [np.percentile(all_times, k, interpolation="nearest") for k in [1, 5, 95, 99]] |
|
if not distributed: |
|
logger.info( |
|
f"{msg}: avg={1.0/avg:.1f} it/s, " |
|
f"p1={percentiles[0]:.2g}s, p5={percentiles[1]:.2g}s, " |
|
f"p95={percentiles[2]:.2g}s, p99={percentiles[3]:.2g}s." |
|
) |
|
return |
|
avg_per_gpu = comm.all_gather(avg) |
|
percentiles_per_gpu = comm.all_gather(percentiles) |
|
if comm.get_rank() > 0: |
|
return |
|
for idx, avg, percentiles in zip(count(), avg_per_gpu, percentiles_per_gpu): |
|
logger.info( |
|
f"GPU{idx} {msg}: avg={1.0/avg:.1f} it/s, " |
|
f"p1={percentiles[0]:.2g}s, p5={percentiles[1]:.2g}s, " |
|
f"p95={percentiles[2]:.2g}s, p99={percentiles[3]:.2g}s." |
|
) |
|
|
|
def benchmark_dataset(self, num_iter, warmup=5): |
|
""" |
|
Benchmark the speed of taking raw samples from the dataset. |
|
""" |
|
|
|
def loader(): |
|
while True: |
|
for k in self.sampler: |
|
yield self.dataset[k] |
|
|
|
self._benchmark(loader(), num_iter, warmup, "Dataset Alone") |
|
|
|
def benchmark_mapper(self, num_iter, warmup=5): |
|
""" |
|
Benchmark the speed of taking raw samples from the dataset and map |
|
them in a single process. |
|
""" |
|
|
|
def loader(): |
|
while True: |
|
for k in self.sampler: |
|
yield self.mapper(self.dataset[k]) |
|
|
|
self._benchmark(loader(), num_iter, warmup, "Single Process Mapper (sec/sample)") |
|
|
|
def benchmark_workers(self, num_iter, warmup=10): |
|
""" |
|
Benchmark the dataloader by tuning num_workers to [0, 1, self.num_workers]. |
|
""" |
|
candidates = [0, 1] |
|
if self.num_workers not in candidates: |
|
candidates.append(self.num_workers) |
|
|
|
dataset = MapDataset(self.dataset, self.mapper) |
|
for n in candidates: |
|
loader = build_batch_data_loader( |
|
dataset, |
|
self.sampler, |
|
self.total_batch_size, |
|
num_workers=n, |
|
) |
|
self._benchmark( |
|
iter(loader), |
|
num_iter * max(n, 1), |
|
warmup * max(n, 1), |
|
f"DataLoader ({n} workers, bs={self.per_gpu_batch_size})", |
|
) |
|
del loader |
|
|
|
def benchmark_IPC(self, num_iter, warmup=10): |
|
""" |
|
Benchmark the dataloader where each worker outputs nothing. This |
|
eliminates the IPC overhead compared to the regular dataloader. |
|
|
|
PyTorch multiprocessing's IPC only optimizes for torch tensors. |
|
Large numpy arrays or other data structure may incur large IPC overhead. |
|
""" |
|
n = self.num_workers |
|
dataset = _EmptyMapDataset(MapDataset(self.dataset, self.mapper)) |
|
loader = build_batch_data_loader( |
|
dataset, self.sampler, self.total_batch_size, num_workers=n |
|
) |
|
self._benchmark( |
|
iter(loader), |
|
num_iter * max(n, 1), |
|
warmup * max(n, 1), |
|
f"DataLoader ({n} workers, bs={self.per_gpu_batch_size}) w/o comm", |
|
) |
|
|
|
def benchmark_distributed(self, num_iter, warmup=10): |
|
""" |
|
Benchmark the dataloader in each distributed worker, and log results of |
|
all workers. This helps understand the final performance as well as |
|
the variances among workers. |
|
|
|
It also prints startup time (first iter) of the dataloader. |
|
""" |
|
gpu = comm.get_world_size() |
|
dataset = MapDataset(self.dataset, self.mapper) |
|
n = self.num_workers |
|
loader = build_batch_data_loader( |
|
dataset, self.sampler, self.total_batch_size, num_workers=n |
|
) |
|
|
|
timer = Timer() |
|
loader = iter(loader) |
|
next(loader) |
|
startup_time = timer.seconds() |
|
logger.info("Dataloader startup time: {:.2f} seconds".format(startup_time)) |
|
|
|
comm.synchronize() |
|
|
|
avg, all_times = self._benchmark(loader, num_iter * max(n, 1), warmup * max(n, 1)) |
|
del loader |
|
self._log_time( |
|
f"DataLoader ({gpu} GPUs x {n} workers, total bs={self.total_batch_size})", |
|
avg, |
|
all_times, |
|
True, |
|
) |
|
|
