| |
| 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, |
| ) |
|
|
