message: [tf.data] Add benchmark for batching sparse tensors.
diff: @@ -18,9 +18,12 @@ from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  +import time
  +
  from absl.testing import parameterized
  import numpy as np
  +from tensorflow.python.client import session
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  @@ -461,5 +464,55 @@ class PaddedBatchDatasetTest(test.TestCase, parameterized.TestCase):
  5, padded_shapes=shape_as_tensor)
  +class BatchDatasetBenchmark(test.Benchmark):
  +
  + def benchmarkBatchSparse(self):
  + non_zeros_per_row_values = [0, 1, 5, 10, 100]
  + batch_size_values = [1, 32, 64, 128, 1024]
  +
  + sparse_placeholder = array_ops.sparse_placeholder(dtype=dtypes.int64)
  + batch_size_placeholder = array_ops.placeholder(dtype=dtypes.int64, shape=[])
  +
  + dataset = dataset_ops.Dataset.from_tensors(sparse_placeholder).repeat(
  + ).batch(batch_size_placeholder)
  + iterator = dataset.make_initializable_iterator()
  + next_element = iterator.get_next()
  +
  + for non_zeros_per_row in non_zeros_per_row_values:
  +
  + sparse_value = sparse_tensor.SparseTensorValue(
  + indices=np.arange(non_zeros_per_row, dtype=np.int64)[:, np.newaxis],
  + values=np.arange(non_zeros_per_row, dtype=np.int64),
  + dense_shape=[1000])
  +
  + for batch_size in batch_size_values:
  +
  + with session.Session() as sess:
  + sess.run(iterator.initializer, feed_dict={
  + sparse_placeholder: sparse_value,
  + batch_size_placeholder: batch_size})
  + # Run five steps to warm up the session caches before taking the
  + # first measurement.
  + for _ in range(5):
  + sess.run(next_element.indices.op)
  + deltas = []
  + for _ in range(100):
  + start = time.time()
  + for _ in range(100):
  + sess.run(next_element.indices.op)
  + end = time.time()
  + deltas.append(end - start)
  +
  + median_wall_time = np.median(deltas) / 100.0
  +
  + print('Batch sparse dataset non-zeros per row: %d batch_size: %d '
  + 'wall time: %f'
  + % (non_zeros_per_row, batch_size, median_wall_time))
  + self.report_benchmark(
  + iters=10000, wall_time=median_wall_time,
  + name='benchmark_batch_sparse_dataset_nnz_%d_batch_size_%d' % (
  + non_zeros_per_row, batch_size))
  +
  +
  if __name__ == '__main__':
  test.main()
  