Function raw_benchmark

tensorflow_datasets/core/utils/benchmark.py:162–237

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) ``` Args: ds: Dataset to benchmark. Can be any iterable. Note: The iterable will be

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

def raw_benchmark(
    ds: Iterable[Any],
    *,
    num_iter: Optional[int] = None,
    batch_size: int = 1,
    detailed_stats: bool = False,
) -> RawBenchmarkResult:
  """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)
  ```

  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
    detailed_stats: Whether to collect detailed statistics such as the time that
      each iteration took.

  Returns:
    raw results.
  """
  try:
    total = len(ds)  # pytype: disable=wrong-arg-types
  except TypeError:
    total = num_iter

  if num_iter is not None:
    total = min(total, num_iter)

  results = []
  actual_num_iter = 0
  first_batch_time = None
  start_time = time.perf_counter_ns()
  end_time = start_time
  for _ in tqdm_utils.tqdm(iter(ds), total=total):
    actual_num_iter += 1
    end_time = time.perf_counter_ns()
    if first_batch_time is None:
      first_batch_time = end_time
    if detailed_stats:
      results.append(end_time)
    if num_iter and actual_num_iter >= num_iter:
      break
  if not actual_num_iter:
    raise ValueError('Cannot benchmark dataset with 0 elements.')

  if num_iter and actual_num_iter < num_iter:
    logging.warning(
        'Number of iterations is shorter than expected ({} vs {})'.format(
            actual_num_iter, num_iter
        )