chore: import upstream snapshot with attribution

tensorflow/python/data/benchmarks/map_benchmark.py

@@ -0,0 +1,194 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Benchmarks for `tf.data.Dataset.map()`."""
+import numpy as np
+
+from tensorflow.python.data.benchmarks import benchmark_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import map_op
+from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import map_fn
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import while_loop
+
+
+class MapBenchmark(benchmark_base.DatasetBenchmarkBase):
+  """Benchmarks for `tf.data.Dataset.map()`."""
+
+  def benchmark_chain_of_maps(self):
+
+    def benchmark_helper(chain_length, fn, use_inter_op_parallelism, label,
+                         benchmark_id):
+      dataset = dataset_ops.Dataset.range(10000)
+      for _ in range(chain_length):
+        dataset = map_op._MapDataset(  # pylint: disable=protected-access
+            dataset, fn, use_inter_op_parallelism=use_inter_op_parallelism)
+      self.run_and_report_benchmark(
+          dataset,
+          num_elements=10000,
+          extras={
+              "model_name": "map.benchmark.%d" % benchmark_id,
+              "parameters": "%d" % chain_length,
+          },
+          name="chain_length_%d%s" % (chain_length, label))
+
+    chain_lengths = [0, 1, 2, 5, 10, 20, 50]
+    for chain_length in chain_lengths:
+      benchmark_helper(
+          chain_length=chain_length,
+          fn=lambda x: x + 1,
+          use_inter_op_parallelism=True,
+          label="",
+          benchmark_id=1)
+      benchmark_helper(
+          chain_length=chain_length,
+          fn=lambda x: x + 1,
+          use_inter_op_parallelism=False,
+          label="_single_threaded",
+          benchmark_id=2)
+      benchmark_helper(
+          chain_length=chain_length,
+          fn=lambda x: x,
+          use_inter_op_parallelism=True,
+          label="_short_circuit",
+          benchmark_id=3)
+
+  def benchmark_map_fan_out(self):
+    fan_outs = [1, 2, 5, 10, 20, 50, 100]
+
+    def benchmark_helper(fan_out, fn, use_inter_op_parallelism, label,
+                         benchmark_id):
+      dataset = dataset_ops.Dataset.from_tensors(
+          tuple(0 for _ in range(fan_out))).repeat(None)
+      dataset = map_op._MapDataset(  # pylint: disable=protected-access
+          dataset, fn, use_inter_op_parallelism=use_inter_op_parallelism)
+      self.run_and_report_benchmark(
+          dataset,
+          num_elements=10000,
+          extras={
+              "model_name": "map.benchmark.%d" % benchmark_id,
+              "parameters": "%d" % fan_out,
+          },
+          name="fan_out_%d%s" % (fan_out, label))
+
+    for fan_out in fan_outs:
+      benchmark_helper(
+          fan_out=fan_out,
+          fn=lambda *xs: [x + 1 for x in xs],
+          use_inter_op_parallelism=True,
+          label="",
+          benchmark_id=4)
+      benchmark_helper(
+          fan_out=fan_out,
+          fn=lambda *xs: [x + 1 for x in xs],
+          use_inter_op_parallelism=False,
+          label="_single_threaded",
+          benchmark_id=5)
+      benchmark_helper(
+          fan_out=fan_out,
+          fn=lambda *xs: xs,
+          use_inter_op_parallelism=True,
+          label="_short_circuit",
+          benchmark_id=6)
+
+  def benchmark_sequential_control_flow(self):
+    dataset = dataset_ops.Dataset.from_tensors(100000)
+
+    def fn(x):
+      i = constant_op.constant(0)
+
+      def body(i, x):
+        return math_ops.add(i, 1), x
+
+      return while_loop.while_loop(math_ops.less, body, [i, x])
+
+    num_elements = 1
+    dataset = dataset.map(fn)
+    self.run_and_report_benchmark(
+        dataset,
+        num_elements=num_elements,
+        extras={
+            "model_name": "map.benchmark.8",
+            "parameters": "%d" % num_elements,
+        },
+        name="sequential_control_flow",
+        apply_default_optimizations=True)
+
+  def benchmark_parallel_control_flow(self):
+    dataset = dataset_ops.Dataset.from_tensors(
+        random_ops.random_uniform([100, 10000000]))
+
+    def fn(x):
+      return map_fn.map_fn(
+          lambda y: y * array_ops.transpose(y), x, parallel_iterations=10)
+
+    num_elements = 1
+    dataset = dataset.map(fn)
+    self.run_and_report_benchmark(
+        dataset,
+        num_elements=1,
+        extras={
+            "model_name": "map.benchmark.9",
+            "parameters": "%d" % num_elements,
+        },
+        name="parallel_control_flow",
+        apply_default_optimizations=True)
+
+  def _benchmark_nested_parallel_map(self, cycle_length, num_parallel_calls):
+    k = 1024 * 1024
+    num_map_elements = 10
+    num_range_elements = 2000
+
+    def g(_):
+      return np.random.rand(50 * k).sum()
+
+    def f(_):
+      return dataset_ops.Dataset.range(num_map_elements).map(
+          g, num_parallel_calls=num_parallel_calls)
+
+    dataset = dataset_ops.Dataset.range(num_range_elements)
+    dataset = dataset.interleave(
+        f, cycle_length=cycle_length, num_parallel_calls=dataset_ops.AUTOTUNE)
+
+    cycle_length_str = ("default"
+                        if cycle_length is None else str(cycle_length))
+    num_parallel_calls_str = ("autotune"
+                              if num_parallel_calls == dataset_ops.AUTOTUNE else
+                              str(num_parallel_calls))
+    map_dataset_str = ("map" if num_parallel_calls is None else
+                       "parallel_map_num_parallel_calls_%s" %
+                       num_parallel_calls_str)
+
+    self.run_and_report_benchmark(
+        dataset,
+        num_elements=num_map_elements * num_range_elements,
+        extras={
+            "model_name": "map.benchmark.10",
+            "parameters": "%s_%s" % (cycle_length_str, num_parallel_calls_str),
+        },
+        name=("%s_cycle_length_%s" % (map_dataset_str, cycle_length_str)))
+
+  def benchmark_nested_parallel_map(self):
+    cycle_lengths = [None, 100]
+    nums_parallel_calls = [None, 1, 10, 100, dataset_ops.AUTOTUNE]
+    for cycle_length in cycle_lengths:
+      for num_parallel_calls in nums_parallel_calls:
+        self._benchmark_nested_parallel_map(cycle_length, num_parallel_calls)
+
+
+if __name__ == "__main__":
+  benchmark_base.test.main()