def benchmark_optimizer(self, optimizer, num_iters):
"""Optimizer benchmark with Bidirectional LSTM model on IMDB data.
Args:
optimizer: The optimizer instance to be benchmarked.
num_iters: The number of iterations to run for performance measurement.
"""
model, train_x, train_y = bidirect_imdb_lstm_config()
metrics, wall_time, extras = benchmark_util.measure_performance(
model,
x=train_x,
y=train_y,
batch_size=512,
optimizer=optimizer,
loss="binary_crossentropy",
metrics=["accuracy"])
name = benchmark_util.get_benchmark_name(self._get_name())
metadata = {
"implementation": name[0],
"model_name": "optimizers",
"parameters": "lstm.512",
}
extras.update(metadata)
self.report_benchmark(
iters=num_iters, wall_time=wall_time, metrics=metrics, extras=extras)
def benchmark_layer_call(self, layer_cls, layer_args, inputs, num_iters):
layer = layer_cls(**_get_layer_args(layer_cls, layer_args))
x = _get_input_data(inputs)
fn = functools.partial(layer, x)
name = benchmark_util.get_benchmark_name(self._get_name())
metadata = {"implementation": name[0] + ".layer.call"}
metadata.update(_get_metadata(name))
self.run_report(fn, num_iters, metadata)
def benchmark_layer_call_with_xla(self, layer_cls, layer_args, inputs,
num_iters):
name = benchmark_util.get_benchmark_name(self._get_name())
# TODO(b/173461426)
if layer_cls is tf.keras.layers.Embedding and name[-1] == "GPU":
return
layer = layer_cls(**_get_layer_args(layer_cls, layer_args))
x = _get_input_data(inputs)
layer.call = tf.function(layer.call, jit_compile=True)
fn = functools.partial(layer, x)
metadata = {"implementation": name[0] + ".layer.call.xla"}
metadata.update(_get_metadata(name))
self.run_report(fn, num_iters, metadata)
def benchmark_layer_call_backward_with_function(self, layer_cls,
layer_args, inputs,
num_iters):
layer = layer_cls(**layer_args)
x = _get_input_data(inputs)
layer.call = tf.function(layer.call)
fn = functools.partial(_layer_call_backward, layer, x)
name = benchmark_util.get_benchmark_name(self._get_name())
metadata = {
"implementation": name[0] + ".layer.call.backward.function"
}
metadata.update(_get_metadata(name))
self.run_report(fn, num_iters, metadata)
def benchmark_layer_call_backward_with_xla(self, layer_cls, layer_args,
inputs, num_iters):
name = benchmark_util.get_benchmark_name(self._get_name())
# TODO(b/153480400)
if layer_cls in [
tf.keras.layers.LSTM, tf.keras.layers.Bidirectional,
tf.keras.layers.ConvLSTM2D, tf.keras.layers.GRU,
tf.keras.layers.RNN, tf.keras.layers.SimpleRNN
]:
return
# TODO(b/173461426)
if layer_cls is tf.keras.layers.Embedding and name[-1] == "GPU":
return
layer = layer_cls(**_get_layer_args(layer_cls, layer_args))
x = _get_input_data(inputs)
layer.call = tf.function(layer.call, jit_compile=True)
fn = functools.partial(_layer_call_backward, layer, x)
metadata = {"implementation": name[0] + ".layer.call.backward.xla"}
metadata.update(_get_metadata(name))
self.run_report(fn, num_iters, metadata)
def test_get_benchmark_name(self):
name = "benchmark_layer_call__Conv2D_small_shape"
expected = ["Conv2D", "small", "shape"]
out = benchmark_util.get_benchmark_name(name)
self.assertAllEqual(out, expected)