text: # @license
  # Copyright 2019 Google LLC. 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.
  # =============================================================================
  
  """Benchmarks for TensorFlow.js tfjs-layers and tfjs-node.
  
  These benchmarks compare the inference and training speed of Keras models of
  varying size and architecture, between Python and browser.
  """
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import argparse
  import functools
  import json
  import os
  import subprocess
  import shutil
  import sys
  import tempfile
  import time
  
  from tensorflow import keras
  import numpy as np
  import tensorflow as tf
  # Comparing TF Eager vs TF.js for a fair comparison.
  if hasattr(tf, 'enable_eager_execution'):
    tf.enable_eager_execution()
  from tensorflow.python.client import device_lib
  import tensorflowjs as tfjs
  
  _FIT_BURNIN_EPOCHS = 1  # How many epochs to call fit() for before timing fit().
  _PREDICT_BURNINS = 20  # How many predict() runs to do before timing predict().
  _PREDICT_RUNS = 30  # How many runs of predict() to average over.
  
  
  def _is_gpu_available():
    devices = device_lib.list_local_devices()
    for device in devices:
      if device.device_type == 'GPU':
        return True
    return False
  
  
  def _get_random_inputs_and_outputs(model, batch_size):
    """Synthesize random inputs and outputs based on the model's specs.
  
    Args:
      model: An instance of keras Model.
      batch_size: Desired batch size.
  
    Returns:
      xs: Synthesized random feature tensor(s).
      ys: Synthesized random target tensor(s).
    """
    input_shapes = [[
        int(d) for d in list(inp.shape[1:])] for inp in model.inputs]
    xs = []
    for in_shape in input_shapes:
      x = np.random.rand(*([batch_size] + in_shape))
      xs.append(x)
    if len(xs) == 1:
      xs = xs[0]
  
    output_shapes = [[
        int(d) for d in list(inp.shape[1:])] for inp in model.outputs]
    ys = []
    for output_shape in output_shapes:
      y = np.random.rand(*([batch_size] + output_shape))
      ys.append(y)
    if len(ys) == 1:
      ys = ys[0]
  
    return xs, ys
  
  
  def benchmark_and_serialize_model(model_name,
                                    description,
                                    model_fn,
                                    input_shape,
                                    target_shape,
                                    optimizer,
                                    loss,
                                    batch_size,
                                    train_epochs,
                                    artifacts_dir,
                                    export_saved_model=False):
    """Benchmark a model's fit() and predict() calls; serialize the model.
  
    Args:
      model_fn: A function that takes two arguments: `input_shape` and
        `target_shape`, and returns a `keras.Model` instance. The model does not
        need to have been compiled.
      input_shape: Input shape as a `list` or `tuple` of integers.
      target_shape: Target shape as a `list` or `tuple` of integers.
      optimizer: The optimizer to use during training.
      loss: The loss function to use during training.
      batch_size: Batch size to use for training.
      train_epochs: Number of training epochs, not including the burn-in epoch(s).
      artifacts_dir: Directory to save the data in. The data includes:
        * topology and weights of the models, in TensorFlow.js format
        * metadata and benchmark information in a file named `data.json`,
          including:
          - the name and description of the model
          - the name of the optimizer used during benchmarking of training
          - loss value
          - the input and output shapes of the model
          - benchmark results from Python Keras.
  
    Returns:
      1. Total fit() time per epoch, averaged over the epochs not including the
         burn-in one.
      2. Average predict() time over all the _PREDICT_RUNS.
    """
    if os.path.isdir(artifacts_dir) and os.listdir(artifacts_dir):
      for rel_name in os.listdir(artifacts_dir):
        abs_path = os.path.join(artifacts_dir, rel_name)
        if os.path.isfile(abs_path):
          os.remove(abs_path)
        else:
          shutil.rmtree(abs_path)
  
    environment_info = _get_python_environment_info()
    task_logs = {}
  
    model = model_fn(input_shape, target_shape)
    if train_epochs:
      model.compile(optimizer=optimizer, loss=loss)
  
    xs, ys = _get_random_inputs_and_outputs(model, batch_size)
  
    # Perform fit() burn-in.
    if train_epochs:
      model.fit(xs, ys, batch_size=batch_size, epochs=_FIT_BURNIN_EPOCHS)
  
    # Time fit().
    if train_epochs:
      train_t_begin = time.time()
      model.fit(xs, ys, batch_size=batch_size, epochs=train_epochs)
      train_t_end = time.time()
  
    # Save data about the model and benchmark results.
    if train_epochs:
      train_time = (train_t_end - train_t_begin) * 1e3 / train_epochs
  
      # Collect and format the data for fit().
      task_logs['fit'] = {  # For schema, see 'ModelTrainingBenchmarkRun` in types.ts.
        'taskType': 'model',
        'modelFormat': 'GraphModel' if export_saved_model else 'LayersModel',
        'modelName': model_name,
        'modelDescription': description,
        'functionName': 'fit',
        'endingTimestampMs': int(time.time() * 1e3),
        'batchSize': batch_size,
        'optimizer': optimizer.__class__.__name__.split('.')[-1],
        'loss': loss,
        'numBenchmarkedIterations': train_epochs,
        'numWarmUpIterations': _FIT_BURNIN_EPOCHS,
        'averageTimeMs': train_time
      }
  
    # Perform predict() burn-in.
    for _ in range(_PREDICT_BURNINS):
      model.predict(xs)
    # Time predict() by averaging.
    predict_ts = []
    for _ in range(_PREDICT_RUNS):
      predict_t_begin = time.time()
      model.predict(xs)
      predict_ts.append((time.time() - predict_t_begin) * 1e3)
  
    if export_saved_model:
      tmp_saved_model_dir = tempfile.mkdtemp()
      tf.compat.v1.keras.experimental.export_saved_model(
          model, tmp_saved_model_dir, serving_only=True)
      subprocess.check_output([
          'tensorflowjs_converter',
          '--input_format', 'tf_saved_model',
          '--output_format', 'tfjs_graph_model',
          '--signature_name', 'serving_default',
          '--saved_model_tags', 'serve',
          tmp_saved_model_dir, artifacts_dir])
      # Clean up the temporary SavedModel directory.
      shutil.rmtree(tmp_saved_model_dir)
    else:
      # Save the model and weights.
      tfjs.converters.save_keras_model(model, artifacts_dir)
  
    # Collect and format the data for predict().
    task_logs['predict'] = {  # For schema, see 'ModelTaskLog` in types.ts.
      'taskType': 'model',
      'modelFormat': 'GraphModel' if export_saved_model else 'LayersModel',
      'modelName': model_name,
      'modelDescription': description,
      'functionName': 'predict',
      'endingTimestampMs': int(time.time() * 1e3),
      'batchSize': batch_size,
      'numBenchmarkedIterations': _PREDICT_RUNS,
      'numWarmUpIterations': _PREDICT_BURNINS,
      'timesMs': predict_ts,
      'averageTimeMs': np.mean(predict_ts)
    }
  
    return task_logs
  
  
  def dense_tiny_model_fn(input_shape, target_shape):
    assert len(target_shape) == 1
    input_layer = keras.Input(input_shape)
    dense_1 = keras.layers.Dense(200, activation='relu')
    dense_2 = keras.layers.Dense(target_shape[0])
    output = dense_2(dense_1(input_layer))
    model = keras.Model(input_layer, output)
    return model
  
  
  def dense_large_model_fn(input_shape, target_shape):
    assert len(target_shape) == 1
    input_layer = keras.Input(input_shape)
    dense_1 = keras.layers.Dense(4000, activation='relu')
    dense_2 = keras.layers.Dense(1000, activation='relu')
    dense_3 = keras.layers.Dense(500, activation='relu')
    dense_4 = keras.layers.Dense(target_shape[0])
    output = dense_4(dense_3(dense_2(dense_1(input_layer))))
    model = keras.Model(input_layer, output)
    return model
  
  
  def convolutional_model_fn(num_filters, input_shape, target_shape):
    """2D convolutional model."""
    kernel_size = 3
    pool_size = 2
    assert len(target_shape) == 1
    num_classes = target_shape[0]
    layers = [
        keras.layers.Conv2D(num_filters, kernel_size,
                            padding='valid',
                            input_shape=input_shape),
        keras.layers.Activation('relu'),
        keras.layers.Conv2D(num_filters, kernel_size),
        keras.layers.Activation('relu'),
        keras.layers.MaxPooling2D(pool_size=pool_size),
        keras.layers.Flatten(),
        keras.layers.Dense(128),
        keras.layers.Activation('relu'),
        keras.layers.Dense(num_classes),
        keras.layers.Activation('softmax')
    ]
    model = keras.models.Sequential(layers)
    return model
  
  
  def mobilenet_v2_model_fn(alpha, input_shape, target_shape):
    """MobileNetV2: A ConvNet from Keras Applications."""
    del input_shape, target_shape  # Unused.
    # `weights=None` leads to random weight initialization and downloadnig
    # of weights.
    model = keras.applications.MobileNetV2(alpha=alpha, weights=None)
    return model
  
  
  def attention_model_fn(input_shape, target_shape):
    """Attention-based translation model."""
    del input_shape, target_shape  # Unused.
    model_json = '{"class_name":"Model","config":{"input_layers":[["input_1",0,0],["s0",0,0],["c0",0,0]],"name":"model_1","layers":[{"class_name":"InputLayer","inbound_nodes":[],"name":"input_1","config":{"dtype":"float32","name":"input_1","sparse":false,"batch_input_shape":[null,30,38]}},{"class_name":"InputLayer","inbound_nodes":[],"name":"s0","config":{"dtype":"float32","name":"s0","sparse":false,"batch_input_shape":[null,64]}},{"class_name":"Bidirectional","inbound_nodes":[[["input_1",0,0,{}]]],"name":"bidirectional_1","config":{"trainable":true,"name":"bidirectional_1","merge_mode":"concat","layer":{"class_name":"LSTM","config":{"stateful":false,"units":32,"activation":"tanh","recurrent_activation":"hard_sigmoid","dropout":0,"recurrent_dropout":0,"use_bias":true,"trainable":true,"recurrent_initializer":{"class_name":"Orthogonal","config":{"seed":null,"gain":1}},"bias_constraint":null,"unroll":false,"kernel_initializer":{"class_name":"VarianceScaling","config":{"seed":null,"distribution":"uniform","mode":"fan_avg","scale":1}},"unit_forget_bias":true,"bias_initializer":{"class_name":"Zeros","config":{}},"kernel_constraint":null,"activity_regularizer":null,"return_sequences":true,"recurrent_constraint":null,"recurrent_regularizer":null,"bias_regularizer":null,"go_backwards":false,"implementation":1,"name":"attLSTM_2","kernel_regularizer":null,"return_state":false}}}},{"class_name":"RepeatVector","inbound_nodes":[[["s0",0,0,{}]],[["attLSTM_1",0,0,{}]],[["attLSTM_1",1,0,{}]],[["attLSTM_1",2,0,{}]],[["attLSTM_1",3,0,{}]],[["attLSTM_1",4,0,{}]],[["attLSTM_1",5,0,{}]],[["attLSTM_1",6,0,{}]],[["attLSTM_1",7,0,{}]],[["attLSTM_1",8,0,{}]]],"name":"repeat_vector_1","config":{"n":30,"trainable":true,"name":"repeat_vector_1"}},{"class_name":"Concatenate","inbound_nodes":[[["bidirectional_1",0,0,{}],["repeat_vector_1",0,0,{}]],[["bidirectional_1",0,0,{}],["repeat_vector_1",1,0,{}]],[["bidirectional_1",0,0,{}],["repeat_vector_1",2,0,{}]],[["bidirectional_1",0,0,{}],["repeat_vector_1",3,0,{}]],[["bidirectional_1",0,0,{}],["repeat_vector_1",4,0,{}]],[["bidirectional_1",0,0,{}],["repeat_vector_1",5,0,{}]],[["bidirectional_1",0,0,{}],["repeat_vector_1",6,0,{}]],[["bidirectional_1",0,0,{}],["repeat_vector_1",7,0,{}]],[["bidirectional_1",0,0,{}],["repeat_vector_1",8,0,{}]],[["bidirectional_1",0,0,{}],["repeat_vector_1",9,0,{}]]],"name":"concatenate_1","config":{"trainable":true,"name":"concatenate_1","axis":-1}},{"class_name":"Dense","inbound_nodes":[[["concatenate_1",0,0,{}]],[["concatenate_1",1,0,{}]],[["concatenate_1",2,0,{}]],[["concatenate_1",3,0,{}]],[["concatenate_1",4,0,{}]],[["concatenate_1",5,0,{}]],[["concatenate_1",6,0,{}]],[["concatenate_1",7,0,{}]],[["concatenate_1",8,0,{}]],[["concatenate_1",9,0,{}]]],"name":"attDense_1","config":{"bias_constraint":null,"kernel_constraint":null,"units":10,"activity_regularizer":null,"use_bias":true,"bias_regularizer":null,"trainable":true,"activation":"tanh","name":"attDense_1","kernel_initializer":{"class_name":"VarianceScaling","config":{"seed":null,"distribution":"uniform","mode":"fan_avg","scale":1}},"kernel_regularizer":null,"bias_initializer":{"class_name":"Zeros","config":{}}}},{"class_name":"Dense","inbound_nodes":[[["attDense_1",0,0,{}]],[["attDense_1",1,0,{}]],[["attDense_1",2,0,{}]],[["attDense_1",3,0,{}]],[["attDense_1",4,0,{}]],[["attDense_1",5,0,{}]],[["attDense_1",6,0,{}]],[["attDense_1",7,0,{}]],[["attDense_1",8,0,{}]],[["attDense_1",9,0,{}]]],"name":"attDense_2","config":{"bias_constraint":null,"kernel_constraint":null,"units":1,"activity_regularizer":null,"use_bias":true,"bias_regularizer":null,"trainable":true,"activation":"relu","name":"attDense_2","kernel_initializer":{"class_name":"VarianceScaling","config":{"seed":null,"distribution":"uniform","mode":"fan_avg","scale":1}},"kernel_regularizer":null,"bias_initializer":{"class_name":"Zeros","config":{}}}},{"class_name":"Activation","inbound_nodes":[[["attDense_2",0,0,{}]],[["attDense_2",1,0,{}]],[["attDense_2",2,0,{}]],[["attDense_2",3,0,{}]],[["attDense_2",4,0,{}]],[["attDense_2",5,0,{}]],[["attDense_2",6,0,{}]],[["attDense_2",7,0,{}]],[["attDense_2",8,0,{}]],[["attDense_2",9,0,{}]]],"name":"attention_weights","config":{"trainable":true,"activation":"softmax","name":"attention_weights"}},{"class_name":"Dot","inbound_nodes":[[["attention_weights",0,0,{}],["bidirectional_1",0,0,{}]],[["attention_weights",1,0,{}],["bidirectional_1",0,0,{}]],[["attention_weights",2,0,{}],["bidirectional_1",0,0,{}]],[["attention_weights",3,0,{}],["bidirectional_1",0,0,{}]],[["attention_weights",4,0,{}],["bidirectional_1",0,0,{}]],[["attention_weights",5,0,{}],["bidirectional_1",0,0,{}]],[["attention_weights",6,0,{}],["bidirectional_1",0,0,{}]],[["attention_weights",7,0,{}],["bidirectional_1",0,0,{}]],[["attention_weights",8,0,{}],["bidirectional_1",0,0,{}]],[["attention_weights",9,0,{}],["bidirectional_1",0,0,{}]]],"name":"dot_1","config":{"trainable":true,"name":"dot_1","normalize":false,"axes":1}},{"class_name":"InputLayer","inbound_nodes":[],"name":"c0","config":{"dtype":"float32","name":"c0","sparse":false,"batch_input_shape":[null,64]}},{"class_name":"LSTM","inbound_nodes":[[["dot_1",0,0,{}],["s0",0,0,{}],["c0",0,0,{}]],[["dot_1",1,0,{}],["attLSTM_1",0,0,{}],["attLSTM_1",0,2,{}]],[["dot_1",2,0,{}],["attLSTM_1",1,0,{}],["attLSTM_1",1,2,{}]],[["dot_1",3,0,{}],["attLSTM_1",2,0,{}],["attLSTM_1",2,2,{}]],[["dot_1",4,0,{}],["attLSTM_1",3,0,{}],["attLSTM_1",3,2,{}]],[["dot_1",5,0,{}],["attLSTM_1",4,0,{}],["attLSTM_1",4,2,{}]],[["dot_1",6,0,{}],["attLSTM_1",5,0,{}],["attLSTM_1",5,2,{}]],[["dot_1",7,0,{}],["attLSTM_1",6,0,{}],["attLSTM_1",6,2,{}]],[["dot_1",8,0,{}],["attLSTM_1",7,0,{}],["attLSTM_1",7,2,{}]],[["dot_1",9,0,{}],["attLSTM_1",8,0,{}],["attLSTM_1",8,2,{}]]],"name":"attLSTM_1","config":{"stateful":false,"units":64,"activation":"tanh","recurrent_activation":"hard_sigmoid","dropout":0,"recurrent_dropout":0,"use_bias":true,"trainable":true,"recurrent_initializer":{"class_name":"Orthogonal","config":{"seed":null,"gain":1}},"bias_constraint":null,"unroll":false,"kernel_initializer":{"class_name":"VarianceScaling","config":{"seed":null,"distribution":"uniform","mode":"fan_avg","scale":1}},"unit_forget_bias":true,"bias_initializer":{"class_name":"Zeros","config":{}},"kernel_constraint":null,"activity_regularizer":null,"return_sequences":false,"recurrent_constraint":null,"recurrent_regularizer":null,"bias_regularizer":null,"go_backwards":false,"implementation":1,"name":"attLSTM_1","kernel_regularizer":null,"return_state":true}},{"class_name":"Dense","inbound_nodes":[[["attLSTM_1",0,0,{}]],[["attLSTM_1",1,0,{}]],[["attLSTM_1",2,0,{}]],[["attLSTM_1",3,0,{}]],[["attLSTM_1",4,0,{}]],[["attLSTM_1",5,0,{}]],[["attLSTM_1",6,0,{}]],[["attLSTM_1",7,0,{}]],[["attLSTM_1",8,0,{}]],[["attLSTM_1",9,0,{}]]],"name":"attDense_3","config":{"bias_constraint":null,"kernel_constraint":null,"units":11,"activity_regularizer":null,"use_bias":true,"bias_regularizer":null,"trainable":true,"activation":"softmax","name":"attDense_3","kernel_initializer":{"class_name":"VarianceScaling","config":{"seed":null,"distribution":"uniform","mode":"fan_avg","scale":1}},"kernel_regularizer":null,"bias_initializer":{"class_name":"Zeros","config":{}}}}],"output_layers":[["attDense_3",0,0],["attDense_3",1,0],["attDense_3",2,0],["attDense_3",3,0],["attDense_3",4,0],["attDense_3",5,0],["attDense_3",6,0],["attDense_3",7,0],["attDense_3",8,0],["attDense_3",9,0]]}}';
    model = keras.models.model_from_json(model_json)
    return model
  
  
  _RNN_TYPE_MAP = {
      'SimpleRNN': keras.layers.SimpleRNN,
      'GRU': keras.layers.GRU,
      'LSTM': keras.layers.LSTM
  }
  
  
  def rnn_model_fn(rnn_type, input_shape, target_shape):
    """Recurrent neural network model."""
    rnnConstructor = _RNN_TYPE_MAP[rnn_type]
    layers = [rnnConstructor(target_shape[0], input_shape=input_shape)]
    model = keras.models.Sequential(layers)
    return model
  
  
  def _get_environment_type():
    return ('python-tensorflow-cuda' if tf.test.gpu_device_name() else
            'python-tensorflow-cpu')
  
  
  def _get_python_environment_info():
    environment_info = {  # For schema, see `PythonEnvironmentInfo` in types.ts.
      'type': _get_environment_type()
    }
  
    try:
      # Disable color from `inxi` command.
      environment_info['cpuInfo'] = tf.compat.as_str(
          subprocess.check_output(['inxi', '-c', '0']))
    except:
      pass
    try:
      environment_info['memInfo'] = tf.compat.as_str(
          subprocess.check_output(['free']))
    except:
      pass
    try:
      environment_info['systemInfo'] = tf.compat.as_str(
          subprocess.check_output(['uname', '-a']))
    except:
      pass
  
    python_ver = sys.version_info
    environment_info['pythonVersion'] = '%d.%d.%d' % (
        python_ver.major, python_ver.minor, python_ver.micro)
    environment_info['tensorflowVersion'] = tf.__version__
    environment_info['kerasVersion'] = keras.__version__
    return environment_info
  
  
  def main():
    environment_info = _get_python_environment_info()
    print('Environment info:')
    print(json.dumps(environment_info, indent=2))
  
    suite_log = dict()  # For schema, see `SuiteLog` in types.ts.
    suite_log['data'] = {}
    suite_log['environmentInfo'] = environment_info
  
    # Dense model.
    optimizer = tf.keras.optimizers.SGD()
    loss = 'mean_squared_error'
    batch_size = 128
    train_epochs = 10
    input_shape = [100]
    target_shape = [1]
    names_fns_and_descriptions = [
        ('dense-tiny',
         dense_tiny_model_fn,
         'Input([%d]);Dense(200);Dense(%d)|%s|%s' %
         (input_shape[0], target_shape[0], optimizer, loss)),
        ('dense-large',
         dense_large_model_fn,
         'Input([%d]);Dense(4000);Dense(1000);Dense(500);Dense(%d)|%s|%s' %
         (input_shape[0], target_shape[0], optimizer, loss))]
  
    for model_name, model_fn, description in names_fns_and_descriptions:
      suite_log['data'][model_name] = (
          benchmark_and_serialize_model(
              model_name,
              description,
              model_fn,
              input_shape,
              target_shape,
              optimizer,
              loss,
              batch_size,
              train_epochs,
              os.path.join(FLAGS.data_root, model_name)))
  
    # dense-tiny and dense-large models as TensorFlow SavedModel (inference only;
    # CPU only).
    # TODO(cais): Make this run on tfjs-node-gpu as well.
    if not _is_gpu_available():
      optimizer = None
      loss = None
      batch_size = 128
      train_epochs = 0
      input_shape = [100]
      target_shape = [1]
      names_fns_and_descriptions = [
          ('dense-tiny_GraphModel',
          dense_tiny_model_fn,
          'Input([%d]);Dense(200);Dense(%d)|%s|%s' %
          (input_shape[0], target_shape[0], optimizer, loss)),
          ('dense-large_GraphModel',
          dense_large_model_fn,
          'Input([%d]);Dense(4000);Dense(1000);Dense(500);Dense(%d)|%s|%s' %
          (input_shape[0], target_shape[0], optimizer, loss))]
  
      for model_name, model_fn, description in names_fns_and_descriptions:
        suite_log['data'][model_name] = (
            benchmark_and_serialize_model(
                model_name,
                description,
                model_fn,
                input_shape,
                target_shape,
                optimizer,
                loss,
                batch_size,
                train_epochs,
                os.path.join(FLAGS.data_root, model_name),
                export_saved_model=True))
  
    # Conv2d models.
    optimizer = tf.keras.optimizers.SGD()
    loss = 'categorical_crossentropy'
    train_epochs = 10
    input_shape = [28, 28, 1]
    target_shape = [10]
    names_fns_and_descriptions = [
        ("convolutional-%dfilters" % num_filters,
         functools.partial(convolutional_model_fn, num_filters),
         'Conv2D(%d,3);Conv2D(%d,3);MaxPooling2D(2);'
         'Flatten();Dense(128);Dense(10)|%s|%s' %
         (num_filters, num_filters, optimizer, loss)) for num_filters in
        (1, 2, 4, 8, 16, 24, 26, 28, 30, 32)]
  
    for model_name, model_fn, description in names_fns_and_descriptions:
      suite_log['data'][model_name] = (
          benchmark_and_serialize_model(
              model_name,
              description,
              model_fn,
              input_shape,
              target_shape,
              optimizer,
              loss,
              batch_size,
              train_epochs,
              os.path.join(FLAGS.data_root, model_name)))
  
    # RNN models.
    # TODO(cais): Restore optimizer after the following
    #   error is resolved:
    # "Error: Cannot evaluate flag 'EPSILON': no evaluation function found."
    # optimizer = tf.keras.optimizers.RMSProp()
    # loss = 'categorical_crossentropy'
    # train_epochs = 10
    optimizer = None
    loss = None
    train_epochs = 0
    input_shape = [20, 20]
    target_shape = [20]
    batch_size = 128
    names_fns_and_descriptions = [
        ("rnn-%s" % rnn_type,
         functools.partial(rnn_model_fn, rnn_type),
         '%s(input_shape=%s, target_shape=%s)|%s|%s' %
         (rnn_type, input_shape, target_shape, optimizer, loss))
        for rnn_type in ('SimpleRNN', 'GRU', 'LSTM')]
  
    for model_name, model_fn, description in names_fns_and_descriptions:
      suite_log['data'][model_name] = (
          benchmark_and_serialize_model(
              model_name,
              description,
              model_fn,
              input_shape,
              target_shape,
              optimizer,
              loss,
              batch_size,
              train_epochs,
              os.path.join(FLAGS.data_root, model_name)))
  
    # MobileNetV2 (inference only).
    input_shape = None  # Determine from the Model object itself.
    target_shape = None  # Determine from the Model object itself.
    batch_size = 8
    train_epochs = 0
    optimizer = None
    loss = None
    names_fns_and_descriptions = [[
        'mobilenet_v2_%.3d' % (alpha * 100),
        functools.partial(mobilenet_v2_model_fn, alpha),
        'mobilenet_v2_%.3d' % (alpha * 100)] for alpha in (0.25, 0.5, 0.75, 1)]
    for model_name, model_fn, description in names_fns_and_descriptions:
      suite_log['data'][model_name] = (
          benchmark_and_serialize_model(
              model_name,
              description,
              model_fn,
              input_shape,
              target_shape,
              optimizer,
              loss,
              batch_size,
              train_epochs,
              os.path.join(FLAGS.data_root, model_name)))
  
    # Attention model (inference only).
    input_shape = None  # Determine from the Model object itself.
    target_shape = None  # Determine from the Model object itself.
    batch_size = 32
    train_epochs = 0
    optimizer = None
    loss = None
    names_fns_and_descriptions = [[
        'attention',
        attention_model_fn,
        'Attention-based translation model: '
        'Function model with bidirectional LSTM layers']]
    for model_name, model_fn, description in names_fns_and_descriptions:
      suite_log['data'][model_name] = (
          benchmark_and_serialize_model(
              model_name,
              description,
              model_fn,
              input_shape,
              target_shape,
              optimizer,
              loss,
              batch_size,
              train_epochs,
              os.path.join(FLAGS.data_root, model_name)))
  
    # MobileNetV2 as TensorFlow SavedModel (inference only; CPU only).
    if not _is_gpu_available():
      # TODO(cais): Make the benchmark run on tensorflow-gpu as well.
      input_shape = None  # Determine from the Model object itself.
      target_shape = None  # Determine from the Model object itself.
      batch_size = 8
      train_epochs = 0
      optimizer = None
      loss = None
      names_fns_and_descriptions = [[
          'mobilenet_v2_%.3d_GraphModel' % (alpha * 100),
          functools.partial(mobilenet_v2_model_fn, alpha),
          'mobilenet_v2_%.3d_GraphModel' % (alpha * 100)]
          for alpha in (0.25, 0.5, 0.75, 1)]
      for model_name, model_fn, description in names_fns_and_descriptions:
        suite_log['data'][model_name] = (
            benchmark_and_serialize_model(
                model_name,
                description,
                model_fn,
                input_shape,
                target_shape,
                optimizer,
                loss,
                batch_size,
                train_epochs,
                os.path.join(FLAGS.data_root, model_name),
                export_saved_model=True))
  
    # TODO(cais): Add fitDataset() calls (i.e., equivalent to fit() with a
    #   tf.data.Dataset object i nPython).
  
    with open(os.path.join(FLAGS.data_root, 'benchmarks.json'), 'wt') as f:
      json.dump(suite_log, f)
  
  
  if __name__ == '__main__':
    parser = argparse.ArgumentParser('Benchmarks demo.')
    parser.add_argument(
        'data_root',
        type=str,
        help='Local path for saving the results of benchmarks.')
    parser.add_argument(
        '--hash_converter', type=str, help='Commit hash of tfjs-conveter')
    parser.add_argument(
        '--hash_core', type=str, help='Commit hash of tfjs-core')
    parser.add_argument(
        '--hash_data', type=str, help='Commit hash of tfjs-data')
    parser.add_argument(
        '--hash_layers', type=str, help='Commit hash of tfjs-layers')
  
    FLAGS, _ = parser.parse_known_args()
    main()
  