1# Copyright 2020 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"""Common utils for benchmark.""" 16from __future__ import absolute_import 17from __future__ import division 18from __future__ import print_function 19 20import timeit 21import numpy as np 22 23import tensorflow as tf 24 25from tensorflow.python.keras.benchmarks import distribution_util 26 27 28def get_benchmark_name(name): 29 """Split the suffix of the benchmark name. 30 31 For example, for the name = 'benchmark_layer_call__Conv2D_small_shape', 32 the return value is ['Conv2D', 'small', 'shape']. 33 34 This is to generate the metadata of the benchmark test. 35 36 Args: 37 name: A string, the benchmark name. 38 39 Returns: 40 A list of strings of the suffix in the benchmark name. 41 """ 42 if '__' not in name or '_' not in name: 43 raise ValueError('The format of the benchmark name is wrong.') 44 return name.split('__')[-1].split('_') 45 46 47def generate_benchmark_params_cpu_gpu(*params_list): 48 """Extend the benchmark names with CPU and GPU suffix. 49 50 Args: 51 *params_list: A list of tuples represents the benchmark parameters. 52 53 Returns: 54 A list of strings with the benchmark name extended with CPU and GPU suffix. 55 """ 56 benchmark_params = [] 57 for params in params_list: 58 benchmark_params.extend([ 59 ((param[0] + '_CPU',) + param[1:]) for param in params 60 ]) 61 benchmark_params.extend([ 62 ((param[0] + '_GPU',) + param[1:]) for param in params 63 ]) 64 return benchmark_params 65 66 67def get_keras_examples_metadata(keras_model, 68 batch_size, 69 impl='.keras.cfit_graph'): 70 return { 71 'model_name': 'keras_examples', 72 'implementation': keras_model + impl, 73 'parameters': 'bs_' + str(batch_size), 74 } 75 76 77class TimerCallBack(tf.keras.callbacks.Callback): 78 """Callback for logging time in each epoch or batch.""" 79 80 def __init__(self): 81 self.times = [] 82 self.timer = timeit.default_timer 83 self.startup_time = timeit.default_timer() 84 self.recorded_startup = False 85 86 def on_epoch_begin(self, e, logs): 87 self.epoch_start_time = self.timer() 88 89 def on_epoch_end(self, e, logs): 90 self.times.append(self.timer() - self.epoch_start_time) 91 92 def on_batch_end(self, e, logs): 93 if not self.recorded_startup: 94 self.startup_time = self.timer() - self.startup_time 95 self.recorded_startup = True 96 97 98def measure_performance(model_fn, 99 x=None, 100 y=None, 101 epochs=2, 102 batch_size=32, 103 run_iters=4, 104 optimizer=None, 105 loss=None, 106 metrics=None, 107 verbose=0, 108 num_gpus=0, 109 distribution_strategy='off'): 110 """Run models and measure the performance. 111 112 Args: 113 model_fn: Model function to be benchmarked. 114 x: Input data. See `x` in the `fit()` method of `keras.Model`. 115 y: Target data. See `y` in the `fit()` method of `keras.Model`. 116 epochs: Integer. Number of epochs to train the model. 117 If unspecified, `epochs` will default to 2. 118 batch_size: Integer. Number of samples per gradient update. If unspecified, 119 `batch_size` will default to 32. 120 run_iters: Integer. Number of iterations to run the performance measurement. 121 If unspecified, `run_iters` will default to 4. 122 optimizer: String (name of optimizer) or optimizer instance. See 123 `tf.keras.optimizers`. 124 loss: String (name of objective function), objective function or 125 `tf.keras.losses.Loss` instance. See `tf.keras.losses`. 126 metrics: Lists of metrics to be evaluated by the model during training. See 127 `metrics` in the `compile()` method of `keras.Model`. 128 verbose: 0, 1, 2. Verbosity mode. See `verbose` in the `fit()` method of 129 `keras.Model`. If unspecified, `verbose` will default to 0. 130 num_gpus: Number of GPUs to run the model. 131 distribution_strategy: Distribution strategies. It could be 132 `multi_worker_mirrored`, `one_device`, `mirrored`. If unspecified, 133 `distribution_strategy` will default to 'off'. Note that, `TPU` 134 and `parameter_server` are not supported yet. 135 136 Returns: 137 Performance summary, which contains build_time, compile_time, 138 startup_time, avg_epoch_time, wall_time, exp_per_sec, epochs, 139 distribution_strategy. 140 141 Raise: 142 ValueError: If `x` is none or if `optimizer` is not provided or 143 if `loss` is not provided or if `num_gpus` is negative. 144 """ 145 if 'x' is None: 146 raise ValueError('Input data is required.') 147 if 'optimizer' is None: 148 raise ValueError('Optimizer is required.') 149 if 'loss' is None: 150 raise ValueError('Loss function is required.') 151 if num_gpus < 0: 152 raise ValueError('`num_gpus` cannot be negative') 153 154 # TODO(xingyulong): we will add tfds support later and 155 # get the `num_examples` from info. 156 num_examples = x.shape[0] 157 158 build_time_list, compile_time_list, startup_time_list = [], [], [] 159 avg_epoch_time_list, wall_time_list, exp_per_sec_list = [], [], [] 160 total_num_examples = epochs * num_examples 161 162 strategy = distribution_util.get_distribution_strategy( 163 distribution_strategy=distribution_strategy, num_gpus=num_gpus) 164 165 for _ in range(run_iters): 166 timer = timeit.default_timer 167 start_time = timer() 168 # Init the distribution strategy scope for each iteration. 169 strategy_scope = distribution_util.get_strategy_scope(strategy) 170 with strategy_scope: 171 t0 = timer() 172 model = model_fn() 173 build_time = timer() - t0 174 175 t1 = timer() 176 model.compile( 177 optimizer=optimizer, 178 loss=loss, 179 metrics=metrics, 180 ) 181 compile_time = timer() - t1 182 # Run one warm up epoch. 183 model.fit(x=x, y=y, batch_size=batch_size, epochs=1) 184 cbk = TimerCallBack() 185 t2 = timer() 186 model.fit( 187 x=x, 188 y=y, 189 batch_size=batch_size, 190 epochs=epochs, 191 callbacks=[cbk], 192 verbose=verbose) 193 end_time = timer() 194 195 build_time_list.append(build_time) 196 compile_time_list.append(compile_time) 197 startup_time_list.append(cbk.startup_time) 198 avg_epoch_time_list.append(np.mean(cbk.times)) 199 wall_time_list.append(end_time - start_time) 200 exp_per_sec_list.append(total_num_examples / (end_time - t2)) 201 202 metrics = [] 203 metrics.append({'name': 'build_time', 'value': np.mean(build_time_list)}) 204 metrics.append({'name': 'compile_time', 'value': np.mean(compile_time_list)}) 205 metrics.append({'name': 'startup_time', 'value': np.mean(startup_time_list)}) 206 metrics.append({ 207 'name': 'avg_epoch_time', 208 'value': np.mean(avg_epoch_time_list) 209 }) 210 metrics.append({'name': 'exp_per_sec', 'value': np.mean(exp_per_sec_list)}) 211 metrics.append({'name': 'epochs', 'value': epochs}) 212 213 wall_time = np.mean(wall_time_list) 214 extras = { 215 'distribution_strategy': distribution_strategy, 216 'num_gpus': num_gpus 217 } 218 219 return metrics, wall_time, extras 220