message: [tf.data] Add two simple microbenchmarks for `Dataset.map()` performance.
diff: @@ -19,13 +19,16 @@ from __future__ import print_function
  from collections import namedtuple
  import threading
  +import time
  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
  from tensorflow.python.framework import errors
  +from tensorflow.python.framework import ops
  from tensorflow.python.framework import sparse_tensor
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import data_flow_ops
  @@ -595,5 +598,64 @@ class MapDatasetTest(test.TestCase):
  with self.assertRaises(errors.OutOfRangeError):
  sess.run(get_next)
  +
  +class MapDatasetBenchmark(test.Benchmark):
  +
  + def benchmarkChainOfMaps(self):
  + chain_lengths = [0, 1, 2, 5, 10, 20, 50]
  + for chain_length in chain_lengths:
  + with ops.Graph().as_default():
  + dataset = dataset_ops.Dataset.from_tensors(0).repeat(None)
  + for _ in range(chain_length):
  + dataset = dataset.map(lambda x: x)
  + iterator = dataset.make_one_shot_iterator()
  + next_element = iterator.get_next()
  +
  + with session.Session() as sess:
  + for _ in range(5):
  + sess.run(next_element.op)
  + deltas = []
  + for _ in range(100):
  + start = time.time()
  + for _ in range(100):
  + sess.run(next_element.op)
  + end = time.time()
  + deltas.append(end - start)
  +
  + median_wall_time = np.median(deltas) / 100
  + print("Map dataset chain length: %d Median wall time: %f"
  + % (chain_length, median_wall_time))
  + self.report_benchmark(
  + iters=1000, wall_time=median_wall_time,
  + name="benchmark_map_dataset_chain_latency_%d" % chain_length)
  +
  + def benchmarkMapFanOut(self):
  + fan_outs = [1, 2, 5, 10, 20, 50, 100]
  + for fan_out in fan_outs:
  + with ops.Graph().as_default():
  + dataset = dataset_ops.Dataset.from_tensors(
  + tuple(0 for _ in range(fan_out))).repeat(None).map(lambda *xs: xs)
  + iterator = dataset.make_one_shot_iterator()
  + next_element = iterator.get_next()
  +
  + with session.Session() as sess:
  + for _ in range(5):
  + sess.run(next_element[0].op)
  + deltas = []
  + for _ in range(100):
  + start = time.time()
  + for _ in range(100):
  + sess.run(next_element[0].op)
  + end = time.time()
  + deltas.append(end - start)
  +
  + median_wall_time = np.median(deltas) / 100
  + print("Map dataset fan out: %d Median wall time: %f"
  + % (fan_out, median_wall_time))
  + self.report_benchmark(
  + iters=1000, wall_time=median_wall_time,
  + name="benchmark_map_dataset_fan_out_%d" % fan_out)
  +
  +
  if __name__ == "__main__":
  test.main()
  