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