ML Community Day is November 9! Join us for updates from TensorFlow, JAX, and more Learn more

tfds.benchmark

Benchmarks any iterable (e.g tf.data.Dataset).

Used in the notebooks

Used in the tutorials

Usage:

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

Reports:

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

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

statistics The recorded statistics, for eventual post-processing