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 17from __future__ import absolute_import 18from __future__ import division 19from __future__ import print_function 20 21import itertools 22import time 23 24from tensorflow.python.framework import ops 25 26 27def 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 34def 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 48def 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