dmfenton/splatt3r

1import json2from collections import defaultdict3from contextlib import contextmanager4from pathlib import Path5from time import time6 7import numpy as np8import torch9import os10 11 12class Benchmarker:13    def __init__(self):14        self.execution_times = defaultdict(list)15 16    @contextmanager17    def time(self, tag: str, num_calls: int = 1):18        try:19            start_time = time()20            yield21        finally:22            end_time = time()23            for _ in range(num_calls):24                self.execution_times[tag].append((end_time - start_time) / num_calls)25 26    def dump(self, path: str) -> None:27        parent = os.path.abspath(os.path.join(path, os.pardir))28        os.makedirs(parent, exist_ok=True)29        # path.parent.mkdir(exist_ok=True, parents=True)30        with open(path, "w") as f:31            json.dump(dict(self.execution_times), f)32        # with path.open("w") as f:33        #     json.dump(dict(self.execution_times), f)34 35    def dump_memory(self, path: Path) -> None:36        path.parent.mkdir(exist_ok=True, parents=True)37        with path.open("w") as f:38            json.dump(torch.cuda.memory_stats()["allocated_bytes.all.peak"], f)39 40    def summarize(self) -> None:41        for tag, times in self.execution_times.items():42            print(f"{tag}: {len(times)} calls, avg. {np.mean(times)} seconds per call")43