| """ | |
| Benchmark the timing a block of code | |
| Authors | |
| * Leo 2022 | |
| """ | |
| import logging | |
| from collections import defaultdict | |
| from contextlib import ContextDecorator | |
| from time import time | |
| from typing import Any | |
| import numpy as np | |
| import torch | |
| logger = logging.getLogger(__name__) | |
| _history = defaultdict(list) | |
| __all__ = ["benchmark"] | |
| class benchmark(ContextDecorator): | |
| def __init__(self, name: str, freq: int = 1) -> None: | |
| super().__init__() | |
| self.name = name | |
| self.freq = freq | |
| def __enter__(self): | |
| torch.cuda.synchronize() | |
| self.start = time() | |
| def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: | |
| torch.cuda.synchronize() | |
| seconds = time() - self.start | |
| global _history | |
| _history[self.name].append(seconds) | |
| if len(_history[self.name]) % self.freq == 0: | |
| logger.warning( | |
| f"{self.name}: {seconds} secs, avg {np.array(_history[self.name]).mean()} secs" | |
| ) | |