# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Benchmark tests for Keras optimizers."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf

from keras.benchmarks import benchmark_util
from keras.optimizer_v2 import adam
from tensorflow.python.platform.benchmark import ParameterizedBenchmark


def bidirect_imdb_lstm_config():
  """Bidirectional LSTM model and IMDB data."""

  def model_fn():
    inputs = tf.keras.Input(shape=(None,), dtype="int32")
    x = tf.keras.layers.Embedding(20000, 128)(inputs)
    x = tf.keras.layers.Bidirectional(
        tf.keras.layers.LSTM(64, return_sequences=True))(
            x)
    x = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(64))(x)
    outputs = tf.keras.layers.Dense(1, activation="sigmoid")(x)
    model = tf.keras.Model(inputs, outputs)
    return model

  (x_train, y_train), _ = tf.keras.datasets.imdb.load_data(num_words=20000)
  x_train = tf.keras.preprocessing.sequence.pad_sequences(x_train, maxlen=200)

  return model_fn, x_train, y_train


class KerasOptimizerBenchmark(
    tf.test.Benchmark, metaclass=ParameterizedBenchmark):
  """Keras optimizer benchmarks."""

  # The parameter of each benchmark test is a tuple, and the first one is
  # the optimizer name.
  _benchmark_parameters = benchmark_util.generate_benchmark_params_cpu_gpu([
      ("Adam", tf.keras.optimizers.Adam(), 10),
      ("NonFusedAdam", adam.NonFusedAdam(), 10),
  ])

  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)


if __name__ == "__main__":
  tf.test.main()

tensorflow.keras.Input tensorflow.keras.layers.Embedding tensorflow.keras.layers.Dense tensorflow.keras.datasets.imdb.load_data tensorflow.test.main tensorflow.keras.Model tensorflow.keras.layers.LSTM tensorflow.keras.optimizers.Adam tensorflow.keras.preprocessing.sequence.pad_sequences
[(41, 'tensorflow.keras.datasets.imdb.load_data', 'tf.keras.datasets.imdb.load_data', ([], {'num_words': '(20000)'}), True, 'import tensorflow as tf\n'), (42, 'tensorflow.keras.preprocessing.sequence.pad_sequences', 'tf.keras.preprocessing.sequence.pad_sequences', (['x_train'], {'maxlen': '(200)'}), True, 'import tensorflow as tf\n'), (86, 'tensorflow.test.main', 'tf.test.main', ([], {}), True, 'import tensorflow as tf\n'), (31, 'tensorflow.keras.Input', 'tf.keras.Input', ([], {'shape': '(None,)', 'dtype': '"""int32"""'}), True, 'import tensorflow as tf\n'), (38, 'tensorflow.keras.Model', 'tf.keras.Model', (['inputs', 'outputs'], {}), True, 'import tensorflow as tf\n'), (66, 'keras.benchmarks.benchmark_util.measure_performance', 'benchmark_util.measure_performance', (['model'], {'x': 'train_x', 'y': 'train_y', 'batch_size': '(512)', 'optimizer': 'optimizer', 'loss': '"""binary_crossentropy"""', 'metrics': "['accuracy']"}), False, 'from keras.benchmarks import benchmark_util\n'), (32, 'tensorflow.keras.layers.Embedding', 'tf.keras.layers.Embedding', (['(20000)', '(128)'], {}), True, 'import tensorflow as tf\n'), (37, 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', (['(1)'], {'activation': '"""sigmoid"""'}), True, 'import tensorflow as tf\n'), (34, 'tensorflow.keras.layers.LSTM', 'tf.keras.layers.LSTM', (['(64)'], {'return_sequences': '(True)'}), True, 'import tensorflow as tf\n'), (36, 'tensorflow.keras.layers.LSTM', 'tf.keras.layers.LSTM', (['(64)'], {}), True, 'import tensorflow as tf\n'), (54, 'tensorflow.keras.optimizers.Adam', 'tf.keras.optimizers.Adam', ([], {}), True, 'import tensorflow as tf\n'), (55, 'keras.optimizer_v2.adam.NonFusedAdam', 'adam.NonFusedAdam', ([], {}), False, 'from keras.optimizer_v2 import adam\n')]