1# Copyright 2015 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"""Simple benchmarks for reductions and their gradients."""
16
17from __future__ import absolute_import
18from __future__ import division
19from __future__ import print_function
20
21import time
22
23import numpy as np
24from six.moves import range  # pylint: disable=redefined-builtin
25
26from tensorflow.core.protobuf import config_pb2
27from tensorflow.python.client import session
28from tensorflow.python.eager import backprop
29from tensorflow.python.eager import context
30from tensorflow.python.framework import constant_op
31from tensorflow.python.framework import ops
32from tensorflow.python.ops import array_ops
33from tensorflow.python.ops import gradients_impl
34from tensorflow.python.ops import math_ops
35from tensorflow.python.platform import test
36
37
38class ReduceBenchmarks(test.Benchmark):
39  """Benchmarks for reductions."""
40
41  def _run(self, func, num_iters):
42    # call func to maybe warm up the GPU
43    func()
44    start = time.time()
45    for _ in range(num_iters):
46      func()
47    end = time.time()
48    mean_us = (end - start) * 1e6 / num_iters
49    self.report_benchmark(
50        iters=num_iters,
51        wall_time=mean_us,
52        extras={"examples_per_sec": num_iters / (end - start)})
53
54  def benchmark_reduce_sum_grad_eager(self):
55    with context.eager_mode():
56      tensor = array_ops.zeros([100, 1000])
57
58      def fn():
59        backprop.gradients_function(math_ops.reduce_sum, [0])(tensor)
60
61      self._run(fn, 10000)
62
63  def benchmark_reduce_sum_grad_eager_cpu(self):
64    with context.eager_mode(), ops.device("/cpu:0"):
65      tensor = array_ops.zeros([100, 1000])
66
67      def fn():
68        backprop.gradients_function(math_ops.reduce_sum, [0])(tensor)
69
70      self._run(fn, 10000)
71
72  def benchmark_reduce_sum_grad_graph(self):
73    config = config_pb2.ConfigProto(
74        graph_options=config_pb2.GraphOptions(
75            optimizer_options=config_pb2.OptimizerOptions(
76                opt_level=config_pb2.OptimizerOptions.L0)))
77    with ops.Graph().as_default(), session.Session(config=config) as sess:
78
79      tensor = constant_op.constant(np.zeros([100, 1000], dtype=np.float32))
80      reduction = math_ops.reduce_sum(tensor)
81      grad, = gradients_impl.gradients(reduction, tensor)
82
83      def fn():
84        self.evaluate(grad.op)
85
86      self._run(fn, 10000)
87
88  def benchmark_reduce_sum_grad_graph_cpu(self):
89    config = config_pb2.ConfigProto(
90        graph_options=config_pb2.GraphOptions(
91            optimizer_options=config_pb2.OptimizerOptions(
92                opt_level=config_pb2.OptimizerOptions.L0)))
93    with ops.Graph().as_default(), session.Session(config=config) as sess:
94
95      with ops.device("/cpu:0"):
96        tensor = constant_op.constant(np.zeros([100, 1000], dtype=np.float32))
97        reduction = math_ops.reduce_sum(tensor)
98        grad, = gradients_impl.gradients(reduction, tensor)
99
100      def fn():
101        self.evaluate(grad.op)
102
103      self._run(fn, 10000)
104
105
106if __name__ == "__main__":
107  test.main()
108