Function benchmark

tensorflow_datasets/core/utils/benchmark.py:240–276

Benchmarks any iterable (e.g `tf.data.Dataset`). Usage: ```py ds = tfds.load('mnist', split='train') ds = ds.batch(32).prefetch(buffer_size=tf.data.AUTOTUNE) tfds.benchmark(ds, batch_size=32) ``` Reports: - Total execution time - Setup time (first warmup batch) - Number of ex

(
    ds: Iterable[Any],
    *,
    num_iter: Optional[int] = None,
    batch_size: int = 1,
)

def benchmark(
    ds: Iterable[Any],
    *,
    num_iter: Optional[int] = None,
    batch_size: int = 1,
) -> BenchmarkResult:
  """Benchmarks any iterable (e.g `tf.data.Dataset`).

  Usage:

  ```py
  ds = tfds.load('mnist', split='train')
  ds = ds.batch(32).prefetch(buffer_size=tf.data.AUTOTUNE)
  tfds.benchmark(ds, batch_size=32)
  ```

  Reports:

  - Total execution time
  - Setup time (first warmup batch)
  - Number of examples/sec

  Args:
    ds: Dataset to benchmark. Can be any iterable. Note: The iterable will be
      fully consumed.
    num_iter: Number of iteration to perform (iteration might be batched)
    batch_size: Batch size of the dataset, used to normalize iterations

  Returns:
    statistics: The recorded statistics, for eventual post-processing
  """
  print('\n************ Summary ************\n')
  raw_results = raw_benchmark(ds=ds, num_iter=num_iter, batch_size=batch_size)
  return BenchmarkResult(
      stats=raw_results.stats_pd(),
      raw_stats=raw_results.raw_stats_pd(),
  )


def _log_stats(msg: str, start: int, end: int, num_examples: int) -> StatDict: