• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Test utilities for tf.data benchmarking functionality."""
16import timeit
17
18import numpy as np
19
20from tensorflow.python.client import session
21from tensorflow.python.data.experimental.ops import testing
22from tensorflow.python.data.ops import dataset_ops
23from tensorflow.python.data.ops import options as options_lib
24from tensorflow.python.eager import context
25from tensorflow.python.platform import test
26
27
28class MetaBenchmark(test.Benchmark):
29  """Benchmark that compares various ways of running tf.data benchmarks."""
30
31  # Note that each of these benchmarks is a separate method so that we can
32  # run them independently and collect a performance profile.
33
34  def setup_fast_dataset(self):
35    self.num_reps = 15
36    self.iters = 100000
37    options = options_lib.Options()
38    options.experimental_optimization.apply_default_optimizations = False
39    return dataset_ops.Dataset.range(10000**2).with_options(options)
40
41  def benchmark_fast_dataset_with_only_cpp_iterations(self):
42    dataset = self.setup_fast_dataset()
43    self.run_benchmark_with_only_cpp_iterations(dataset)
44
45  def benchmark_fast_dataset_with_session_run(self):
46    dataset = self.setup_fast_dataset()
47    self.run_benchmark_with_session_run(dataset)
48
49  def benchmark_fast_dataset_with_session_callable(self):
50    dataset = self.setup_fast_dataset()
51    self.run_benchmark_with_session_run(dataset, make_callable=True)
52
53  def benchmark_fast_dataset_in_eager(self):
54    with context.eager_mode():
55      dataset = self.setup_fast_dataset()
56      self.run_benchmark_in_eager(dataset)
57
58  def setup_slow_dataset(self):
59    dataset = self.setup_fast_dataset()
60    self.iters = 1000
61    # sleep for 1e-3s per iteration
62    return dataset.apply(testing.sleep(1000))
63
64  def benchmark_slow_dataset_with_only_cpp_iterations(self):
65    dataset = self.setup_slow_dataset()
66    self.run_benchmark_with_only_cpp_iterations(dataset)
67
68  def benchmark_slow_dataset_with_session_run(self):
69    dataset = self.setup_slow_dataset()
70    self.run_benchmark_with_session_run(dataset)
71
72  def benchmark_slow_dataset_with_session_callable(self):
73    dataset = self.setup_slow_dataset()
74    self.run_benchmark_with_session_run(dataset, make_callable=True)
75
76  def benchmark_slow_dataset_in_eager(self):
77    with context.eager_mode():
78      dataset = self.setup_slow_dataset()
79      self.run_benchmark_in_eager(dataset)
80
81  def report(self, deltas):
82    # Each `delta` is the time taken for `self.iters` iterations. Divide by the
83    # number of iterations here to get per-element iteration time.
84    deltas = np.array(deltas) / self.iters
85    # Discard the first 5 results from "warming up" the session.
86    deltas = deltas[5:]
87
88    median = np.median(deltas)
89    mean = np.mean(deltas)
90    min_val = np.min(deltas)
91    max_val = np.max(deltas)
92    extras = {
93        "iters_per_second": 1 / median,
94        "median": median,
95        "mean": mean,
96        "min": min_val,
97        "max": max_val,
98        "num_reps": self.num_reps - 5,
99    }
100    self.report_benchmark(wall_time=median, iters=self.iters, extras=extras)
101
102  def run_benchmark_in_eager(self, dataset):
103    deltas = []
104    for _ in range(self.num_reps):
105      iterator = iter(dataset)
106      deltas.append(timeit.timeit(lambda: next(iterator), number=self.iters))  # pylint: disable=cell-var-from-loop
107
108    self.report(deltas)
109
110  def run_benchmark_with_session_run(self, dataset, make_callable=False):
111    iterator = dataset_ops.make_initializable_iterator(dataset)
112    next_element = iterator.get_next()
113
114    with session.Session() as sess:
115      deltas = []
116      for _ in range(self.num_reps):
117        if make_callable:
118          get_next_element = sess.make_callable(next_element)
119        else:
120          # Note: session.run(next_element.op) is more performant than
121          # session.run(next_element) because we avoid the cost of copying the
122          # tensor from C++ to python.
123          get_next_element = lambda: sess.run(next_element.op)
124
125        sess.run(iterator.initializer)
126        deltas.append(timeit.timeit(get_next_element, number=self.iters))
127    self.report(deltas)
128
129  def run_benchmark_with_only_cpp_iterations(self, dataset):
130    """Benchmarks the dataset with the iterations performed in C++."""
131    # NOTE: We use `dataset.skip()` to perform the iterations in C++, avoiding
132    # the overhead of multiple `session.run()` calls. Note that this relies on
133    # the underlying implementation of `skip`: if it is optimized in the future,
134    # we will have to change this code.
135    dataset = dataset.skip(self.iters - 1)
136    iterator = dataset_ops.make_initializable_iterator(dataset)
137    next_element = iterator.get_next()
138
139    with session.Session() as sess:
140      deltas = []
141      for _ in range(self.num_reps):
142        sess.run(iterator.initializer)
143        deltas.append(
144            timeit.timeit(lambda: sess.run(next_element.op), number=1))
145    self.report(deltas)
146
147
148if __name__ == "__main__":
149  test.main()
150