experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Simplify layer benchmarks.

Browse Source 
PiperOrigin-RevId: 351672917
Change-Id: I288ca4e2e7e78912d2d1d826e9942b4f0474068e
This commit is contained in:
Yanhui Liang 2021-01-13 15:04:14 -08:00 committed by TensorFlower Gardener
parent 42e87dcac5
commit 147dd154b4
1 changed files with 9 additions and 47 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

56
tensorflow/python/keras/benchmarks/layer_benchmarks/layer_benchmarks_test.py
View File

@ -68,38 +68,38 @@ class KerasLayerBenchmarks(six.with_metaclass(
"activation": "relu"
}, {
"input_shape": (1, 1, 1, 1)
}, 10),
}, 100),
("Conv2D_normal_shape", tf.keras.layers.Conv2D, {
"filters": 1,
"kernel_size": 1,
"activation": "relu"
}, {
"input_shape": (64, 28, 28, 3)
}, 10),
}, 100),
("LSTM_small_shape", tf.keras.layers.LSTM, {
"units": 1
}, {
"input_shape": (1, 1, 1)
}, 10),
}, 100),
("LSTM_normal_shape", tf.keras.layers.LSTM, {
"units": 4
}, {
"input_shape": (32, 10, 8)
}, 10),
}, 100),
("Embedding_small_shape", tf.keras.layers.Embedding, {
"input_dim": 1,
"output_dim": 1,
"input_length": 1
}, {
"input": np.random.randint(1, size=(1, 1))
}, 10),
}, 100),
("Embedding_normal_shape", tf.keras.layers.Embedding, {
"input_dim": 1000,
"output_dim": 64,
"input_length": 10
}, {
"input": np.random.randint(1000, size=(32, 10))
}, 10),
}, 100),
])
def benchmark_layer_call(self, layer_cls, layer_args, inputs, num_iters):
@ -163,50 +163,12 @@ class KerasLayerBenchmarks(six.with_metaclass(
metadata.update(_get_metadata(name))
self.run_report(fn, num_iters, metadata)
class KerasLayerBenchmarksBackwardXLA(six.with_metaclass(
benchmark.ParameterizedBenchmark,
layer_benchmarks_test_base.LayerBenchmarksBase)):
_benchmark_parameters = benchmark_util.generate_benchmark_params_cpu_gpu([
("Conv2D_small_shape", tf.keras.layers.Conv2D, {
"filters": 1,
"kernel_size": 1,
"activation": "relu"
}, {
"input_shape": (1, 1, 1, 1)
}, 10000),
("Conv2D_normal_shape", tf.keras.layers.Conv2D, {
"filters": 1,
"kernel_size": 1,
"activation": "relu"
}, {
"input_shape": (64, 28, 28, 3)
}, 10000),
# TODO(b/153480400)
# ("LSTM_small_shape", tf.keras.layers.LSTM,
# {"units": 1}, {"input_shape": (1, 1, 1)}, 10000),
# ("LSTM_normal_shape", tf.keras.layers.LSTM,
# {"units": 4}, {"input_shape": (32, 10, 8)}, 10000),
("Embedding_small_shape", tf.keras.layers.Embedding, {
"input_dim": 1,
"output_dim": 1,
"input_length": 1
}, {
"input": np.random.randint(1, size=(1, 1))
}, 10),
("Embedding_normal_shape", tf.keras.layers.Embedding, {
"input_dim": 1000,
"output_dim": 64,
"input_length": 10
}, {
"input": np.random.randint(1000, size=(32, 10))
}, 10),
])
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 is tf.keras.layers.LSTM:
return
# TODO(b/173461426)
if layer_cls is tf.keras.layers.Embedding and name[-1] == "GPU":
return