tfds.core.benchmark

View source on GitHub

Benchmarks a tf.data.Dataset.

tfds.core.benchmark(
    ds: tf.data.Dataset,
    num_iter: Optional[int] = None,
    batch_size: int = 1
) -> Dict[str, StatDict]

Usage:

ds = tfds.load('mnist', split='train').batch(32).prefetch()
tfds.core.benchmark(ds, batch_size=32)

Reports:

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

ds Dataset to benchmark
num_iter Number of iteration to perform (iteration might be batched)
batch_size Batch size of the dataset, used to normalize iterations

statistics The recorded statistics, for eventual post-processing