diff --git a/experanto/utils.py b/experanto/utils.py index d2891320..fdb8e145 100644 --- a/experanto/utils.py +++ b/experanto/utils.py @@ -121,6 +121,42 @@ def __iter__(self): for i in range(len(self)): yield next(self.iterator) + def benchmark( + self, + n_batches=100, + n_warmup=10, + verbose=False, + ): + it = iter(self) + batch_times = [] + + total_batches = len(self) if n_batches is None else n_batches + + # Warmup + for _ in range(n_warmup): + batch = next(it, None) + if batch is None: + it = iter(self) + batch = next(it) + + # Timed iteration + for i in range(total_batches): + start = time.perf_counter() + batch = next(it, None) + if batch is None: + it = iter(self) + batch = next(it) + end = time.perf_counter() + batch_times.append(end - start) + if verbose: + print(f"Batch {i+1}: {batch_times[-1]:.4f}s") + + avg_time = sum(batch_times) / len(batch_times) + std_time = ( + sum((t - avg_time) ** 2 for t in batch_times) / len(batch_times) + ) ** 0.5 + return {"avg_time": avg_time, "std_time": std_time, "batch_times": batch_times} + # borrowed with <3 from # https://github.com/sinzlab/neuralpredictors/blob/main/neuralpredictors/training/cyclers.py