1 # Copyright 2017 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 """Library for benchmarking OpKernels.""" 16 17 from __future__ import absolute_import 18 from __future__ import division 19 from __future__ import print_function 20 21 import itertools 22 import time 23 24 from tensorflow.python.framework import ops 25 26 27 def device(use_gpu=False): 28 """TensorFlow device to assign ops to.""" 29 if use_gpu: 30 return ops.device("/gpu:0") 31 return ops.device("/cpu:0") 32 33 34 def seconds_per_run(op, sess, num_runs=50): 35 """Number of seconds taken to execute 'op' once on average.""" 36 for _ in range(2): 37 sess.run(op) 38 39 start_time = time.time() 40 for _ in range(num_runs): 41 sess.run(op) 42 43 end_time = time.time() 44 time_taken = (end_time - start_time) / num_runs 45 return time_taken 46 47 48 def dict_product(dicts): 49 """Constructs iterator over outer product of entries in a dict-of-lists. 50 51 Example: 52 >>> dict_products({"a": [1,2], "b": [3, 4]}) 53 >>> [{"a": 1, "b": 3}, 54 {"a": 1, "b": 4}, 55 {"a": 2, "b": 3}, 56 {"a": 2, "b": 4}] 57 58 Args: 59 dicts: dictionary with string keys and list values. 60 61 Yields: 62 Individual dicts from outer product. 63 """ 64 keys, values = zip(*dicts.items()) 65 for config_values in itertools.product(*values): 66 yield dict(zip(keys, config_values)) 67
