diff --git a/tensorflow/python/keras/benchmarks/layer_benchmarks/layer_benchmarks_test.py b/tensorflow/python/keras/benchmarks/layer_benchmarks/layer_benchmarks_test.py index 850e7efdf45..2a7bede8fcb 100644 --- a/tensorflow/python/keras/benchmarks/layer_benchmarks/layer_benchmarks_test.py +++ b/tensorflow/python/keras/benchmarks/layer_benchmarks/layer_benchmarks_test.py @@ -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