# Copyright 2018 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.
# ==============================================================================
# pylint: disable=protected-access
"""Tests for saving/loading function for keras Model."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import shutil

from absl.testing import parameterized
import numpy as np

from tensorflow.python import keras
from tensorflow.python import tf2
from tensorflow.python.client import session
from tensorflow.python.eager import context
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_spec
from tensorflow.python.framework import test_util
from tensorflow.python.keras import keras_parameterized
from tensorflow.python.keras import testing_utils
from tensorflow.python.keras.engine import training as model_lib
from tensorflow.python.keras.optimizer_v2 import adadelta
from tensorflow.python.keras.optimizer_v2 import rmsprop
from tensorflow.python.keras.saving import saved_model_experimental as keras_saved_model
from tensorflow.python.keras.utils import mode_keys
from tensorflow.python.keras.utils import tf_utils
from tensorflow.python.ops import array_ops
from tensorflow.python.platform import test
from tensorflow.python.saved_model import loader_impl
from tensorflow.python.saved_model import model_utils
from tensorflow.python.training import training as training_module


@keras_parameterized.run_all_keras_modes()
class TestModelSavingandLoading(parameterized.TestCase, test.TestCase):

  def _save_model_dir(self, dirname='saved_model'):
    temp_dir = self.get_temp_dir()
    self.addCleanup(shutil.rmtree, temp_dir, ignore_errors=True)
    return os.path.join(temp_dir, dirname)

  def test_saving_sequential_model(self):
    with self.cached_session():
      model = keras.models.Sequential()
      model.add(keras.layers.Dense(2, input_shape=(3,)))
      model.add(keras.layers.RepeatVector(3))
      model.add(keras.layers.TimeDistributed(keras.layers.Dense(3)))
      model.compile(
          loss=keras.losses.MSE,
          optimizer=rmsprop.RMSprop(lr=0.0001),
          metrics=[keras.metrics.categorical_accuracy],
          sample_weight_mode='temporal',
          run_eagerly=testing_utils.should_run_eagerly(),
          experimental_run_tf_function=testing_utils.should_run_tf_function())
      x = np.random.random((1, 3))
      y = np.random.random((1, 3, 3))
      model.train_on_batch(x, y)

      ref_y = model.predict(x)

      saved_model_dir = self._save_model_dir()
      keras_saved_model.export_saved_model(model, saved_model_dir)

      loaded_model = keras_saved_model.load_from_saved_model(saved_model_dir)
      y = loaded_model.predict(x)
      self.assertAllClose(ref_y, y, atol=1e-05)

  @test_util.run_in_graph_and_eager_modes
  def test_saving_sequential_model_without_compile(self):
    with self.cached_session():
      model = keras.models.Sequential()
      model.add(keras.layers.Dense(2, input_shape=(3,)))
      model.add(keras.layers.RepeatVector(3))
      model.add(keras.layers.TimeDistributed(keras.layers.Dense(3)))

      x = np.random.random((1, 3))
      ref_y = model.predict(x)

      saved_model_dir = self._save_model_dir()
      keras_saved_model.export_saved_model(model, saved_model_dir)
      loaded_model = keras_saved_model.load_from_saved_model(saved_model_dir)

      y = loaded_model.predict(x)
      self.assertAllClose(ref_y, y, atol=1e-05)

  def test_saving_functional_model(self):
    with self.cached_session():
      inputs = keras.layers.Input(shape=(3,))
      x = keras.layers.Dense(2)(inputs)
      output = keras.layers.Dense(3)(x)

      model = keras.models.Model(inputs, output)
      model.compile(
          loss=keras.losses.MSE,
          optimizer=rmsprop.RMSprop(lr=0.0001),
          metrics=[keras.metrics.categorical_accuracy],
          run_eagerly=testing_utils.should_run_eagerly(),
          experimental_run_tf_function=testing_utils.should_run_tf_function())
      x = np.random.random((1, 3))
      y = np.random.random((1, 3))
      model.train_on_batch(x, y)

      ref_y = model.predict(x)

      saved_model_dir = self._save_model_dir()
      keras_saved_model.export_saved_model(model, saved_model_dir)
      loaded_model = keras_saved_model.load_from_saved_model(saved_model_dir)

      y = loaded_model.predict(x)
      self.assertAllClose(ref_y, y, atol=1e-05)

  @test_util.run_in_graph_and_eager_modes
  def test_saving_functional_model_without_compile(self):
    with self.cached_session():
      inputs = keras.layers.Input(shape=(3,))
      x = keras.layers.Dense(2)(inputs)
      output = keras.layers.Dense(3)(x)

      model = keras.models.Model(inputs, output)

      x = np.random.random((1, 3))
      y = np.random.random((1, 3))

      ref_y = model.predict(x)

      saved_model_dir = self._save_model_dir()
      keras_saved_model.export_saved_model(model, saved_model_dir)
      loaded_model = keras_saved_model.load_from_saved_model(saved_model_dir)

      y = loaded_model.predict(x)
      self.assertAllClose(ref_y, y, atol=1e-05)

  @test_util.run_in_graph_and_eager_modes
  def test_saving_with_tf_optimizer(self):
    model = keras.models.Sequential()
    model.add(keras.layers.Dense(2, input_shape=(3,)))
    model.add(keras.layers.Dense(3))
    model.compile(
        loss='mse',
        optimizer=training_module.RMSPropOptimizer(0.1),
        metrics=['acc'])

    x = np.random.random((1, 3))
    y = np.random.random((1, 3))
    model.train_on_batch(x, y)
    ref_y = model.predict(x)

    saved_model_dir = self._save_model_dir()
    keras_saved_model.export_saved_model(model, saved_model_dir)
    loaded_model = keras_saved_model.load_from_saved_model(saved_model_dir)
    loaded_model.compile(
        loss='mse',
        optimizer=training_module.RMSPropOptimizer(0.1),
        metrics=['acc'],
        run_eagerly=testing_utils.should_run_eagerly(),
        experimental_run_tf_function=testing_utils.should_run_tf_function())
    y = loaded_model.predict(x)
    self.assertAllClose(ref_y, y, atol=1e-05)

    # test that new updates are the same with both models
    x = np.random.random((1, 3))
    y = np.random.random((1, 3))

    ref_loss = model.train_on_batch(x, y)
    loss = loaded_model.train_on_batch(x, y)
    self.assertAllClose(ref_loss, loss, atol=1e-05)

    ref_y = model.predict(x)
    y = loaded_model.predict(x)
    self.assertAllClose(ref_y, y, atol=1e-05)

    # test saving/loading again
    saved_model_dir2 = self._save_model_dir('saved_model_2')
    keras_saved_model.export_saved_model(loaded_model, saved_model_dir2)
    loaded_model = keras_saved_model.load_from_saved_model(saved_model_dir2)
    y = loaded_model.predict(x)
    self.assertAllClose(ref_y, y, atol=1e-05)

  def test_saving_subclassed_model_raise_error(self):
    # For now, saving subclassed model should raise an error. It should be
    # avoided later with loading from SavedModel.pb.

    class SubclassedModel(model_lib.Model):

      def __init__(self):
        super(SubclassedModel, self).__init__()
        self.layer1 = keras.layers.Dense(3)
        self.layer2 = keras.layers.Dense(1)

      def call(self, inp):
        return self.layer2(self.layer1(inp))

    model = SubclassedModel()

    saved_model_dir = self._save_model_dir()
    with self.assertRaises(NotImplementedError):
      keras_saved_model.export_saved_model(model, saved_model_dir)


class LayerWithLearningPhase(keras.engine.base_layer.Layer):

  def build(self, input_shape):
    self.input_spec = keras.layers.InputSpec(shape=[None] * len(input_shape))
    self.built = True

  def call(self, x, training=None):
    if training is None:
      training = keras.backend.learning_phase()
    output = tf_utils.smart_cond(
        training, lambda: x * 0, lambda: array_ops.identity(x))
    if not context.executing_eagerly():
      output._uses_learning_phase = True  # pylint: disable=protected-access
    return output

  def compute_output_shape(self, input_shape):
    return input_shape


def functional_model(uses_learning_phase=True):
  inputs = keras.layers.Input(shape=(3,))
  x = keras.layers.Dense(2)(inputs)
  x = keras.layers.Dense(3)(x)
  if uses_learning_phase:
    x = LayerWithLearningPhase()(x)
  return keras.models.Model(inputs, x)


def sequential_model(uses_learning_phase=True):
  model = keras.models.Sequential()
  model.add(keras.layers.Dense(2, input_shape=(3,)))
  model.add(keras.layers.Dense(3))
  if uses_learning_phase:
    model.add(LayerWithLearningPhase())
  return model


def sequential_model_without_input_shape(uses_learning_phase=True):
  model = keras.models.Sequential()
  model.add(keras.layers.Dense(2))
  model.add(keras.layers.Dense(3))
  if uses_learning_phase:
    model.add(LayerWithLearningPhase())
  return model


class Subclassed(keras.models.Model):

  def __init__(self):
    super(Subclassed, self).__init__()
    self.dense1 = keras.layers.Dense(2)
    self.dense2 = keras.layers.Dense(3)

  def call(self, inputs):
    x = self.dense1(inputs)
    x = self.dense2(x)
    return x


def subclassed_model():
  return Subclassed()


def load_model(sess, path, mode):
  tags = model_utils.EXPORT_TAG_MAP[mode]
  sig_def_key = model_utils.SIGNATURE_KEY_MAP[mode]

  meta_graph_def = loader_impl.load(sess, tags, path)
  inputs = {
      k: sess.graph.get_tensor_by_name(v.name)
      for k, v in meta_graph_def.signature_def[sig_def_key].inputs.items()}
  outputs = {
      k: sess.graph.get_tensor_by_name(v.name)
      for k, v in meta_graph_def.signature_def[sig_def_key].outputs.items()}
  return inputs, outputs, meta_graph_def


@test_util.run_all_in_graph_and_eager_modes
class TestModelSavedModelExport(test.TestCase, parameterized.TestCase):

  def _save_model_dir(self, dirname='saved_model'):
    temp_dir = self.get_temp_dir()
    self.addCleanup(shutil.rmtree, temp_dir, ignore_errors=True)
    return os.path.join(temp_dir, dirname)

  @parameterized.parameters(
      {
          'model_builder': functional_model,
          'uses_learning_phase': True,
          'optimizer_cls': adadelta.Adadelta,
          'train_before_export': True},
      {
          'model_builder': functional_model,
          'uses_learning_phase': True,
          'optimizer_cls': training_module.AdadeltaOptimizer,
          'train_before_export': False},
      {
          'model_builder': functional_model,
          'uses_learning_phase': False,
          'optimizer_cls': None,
          'train_before_export': False},
      {
          'model_builder': sequential_model,
          'uses_learning_phase': True,
          'optimizer_cls': training_module.AdadeltaOptimizer,
          'train_before_export': True},
      {
          'model_builder': sequential_model,
          'uses_learning_phase': True,
          'optimizer_cls': adadelta.Adadelta,
          'train_before_export': False},
      {
          'model_builder': sequential_model,
          'uses_learning_phase': False,
          'optimizer_cls': None,
          'train_before_export': False},
      {
          'model_builder': sequential_model_without_input_shape,
          'uses_learning_phase': True,
          'optimizer_cls': training_module.AdadeltaOptimizer,
          'train_before_export': False})
  def testSaveAndLoadSavedModelExport(
      self, model_builder, uses_learning_phase, optimizer_cls,
      train_before_export):
    optimizer = None if optimizer_cls is None else optimizer_cls()

    saved_model_dir = self._save_model_dir()

    np.random.seed(130)
    input_arr = np.random.random((1, 3))
    target_arr = np.random.random((1, 3))

    model = model_builder(uses_learning_phase)
    if optimizer is not None:
      model.compile(
          loss='mse',
          optimizer=optimizer,
          metrics=['mae'])
      if train_before_export:
        model.train_on_batch(input_arr, target_arr)

      ref_loss, ref_mae = model.evaluate(input_arr, target_arr)

    ref_predict = model.predict(input_arr)

    # Export SavedModel
    keras_saved_model.export_saved_model(model, saved_model_dir)

    input_name = model.input_names[0]
    output_name = model.output_names[0]
    target_name = output_name + '_target'

    # Load predict graph, and test predictions
    with session.Session(graph=ops.Graph()) as sess:
      inputs, outputs, _ = load_model(sess, saved_model_dir,
                                      mode_keys.ModeKeys.PREDICT)

      predictions = sess.run(outputs[output_name],
                             {inputs[input_name]: input_arr})
      self.assertAllClose(ref_predict, predictions, atol=1e-05)

    if optimizer:
      # Load eval graph, and test predictions, loss and metric values
      with session.Session(graph=ops.Graph()) as sess:
        inputs, outputs, _ = load_model(sess, saved_model_dir,
                                        mode_keys.ModeKeys.TEST)

        # First obtain the loss and predictions, and run the metric update op by
        # feeding in the inputs and targets.
        metrics_name = 'mae' if tf2.enabled() else 'mean_absolute_error'
        metrics_update_op_key = 'metrics/' + metrics_name + '/update_op'
        metrics_value_op_key = 'metrics/' + metrics_name + '/value'

        loss, predictions, _ = sess.run(
            (outputs['loss'], outputs['predictions/' + output_name],
             outputs[metrics_update_op_key]), {
                 inputs[input_name]: input_arr,
                 inputs[target_name]: target_arr
             })

        # The metric value should be run after the update op, to ensure that it
        # reflects the correct value.
        metric_value = sess.run(outputs[metrics_value_op_key])

        self.assertEqual(int(train_before_export),
                         sess.run(training_module.get_global_step()))
        self.assertAllClose(ref_loss, loss, atol=1e-05)
        self.assertAllClose(ref_mae, metric_value, atol=1e-05)
        self.assertAllClose(ref_predict, predictions, atol=1e-05)

      # Load train graph, and check for the train op, and prediction values
      with session.Session(graph=ops.Graph()) as sess:
        inputs, outputs, meta_graph_def = load_model(
            sess, saved_model_dir, mode_keys.ModeKeys.TRAIN)
        self.assertEqual(int(train_before_export),
                         sess.run(training_module.get_global_step()))
        self.assertIn('loss', outputs)
        self.assertIn(metrics_update_op_key, outputs)
        self.assertIn(metrics_value_op_key, outputs)
        self.assertIn('predictions/' + output_name, outputs)

        # Train for a step
        train_op = loader_impl.get_train_op(meta_graph_def)
        train_outputs, _ = sess.run(
            [outputs, train_op], {inputs[input_name]: input_arr,
                                  inputs[target_name]: target_arr})
        self.assertEqual(int(train_before_export) + 1,
                         sess.run(training_module.get_global_step()))

        if uses_learning_phase:
          self.assertAllClose(
              [[0, 0, 0]], train_outputs['predictions/' + output_name],
              atol=1e-05)
        else:
          self.assertNotAllClose(
              [[0, 0, 0]], train_outputs['predictions/' + output_name],
              atol=1e-05)

  def testSaveAndLoadSavedModelWithCustomObject(self):
    saved_model_dir = self._save_model_dir()
    with session.Session(graph=ops.Graph()) as sess:
      def relu6(x):
        return keras.backend.relu(x, max_value=6)
      inputs = keras.layers.Input(shape=(1,))
      outputs = keras.layers.Activation(relu6)(inputs)
      model = keras.models.Model(inputs, outputs)
      keras_saved_model.export_saved_model(
          model, saved_model_dir, custom_objects={'relu6': relu6})
    with session.Session(graph=ops.Graph()) as sess:
      inputs, outputs, _ = load_model(sess, saved_model_dir,
                                      mode_keys.ModeKeys.PREDICT)
      input_name = model.input_names[0]
      output_name = model.output_names[0]
      predictions = sess.run(
          outputs[output_name], {inputs[input_name]: [[7], [-3], [4]]})
      self.assertAllEqual([[6], [0], [4]], predictions)

  def testAssertModelCloneSameObjectsIgnoreOptimizer(self):
    input_arr = np.random.random((1, 3))
    target_arr = np.random.random((1, 3))

    model_graph = ops.Graph()
    clone_graph = ops.Graph()

    # Create two models with the same layers but different optimizers.
    with session.Session(graph=model_graph):
      inputs = keras.layers.Input(shape=(3,))
      x = keras.layers.Dense(2)(inputs)
      x = keras.layers.Dense(3)(x)
      model = keras.models.Model(inputs, x)

      model.compile(loss='mse', optimizer=training_module.AdadeltaOptimizer())
      model.train_on_batch(input_arr, target_arr)

    with session.Session(graph=clone_graph):
      inputs = keras.layers.Input(shape=(3,))
      x = keras.layers.Dense(2)(inputs)
      x = keras.layers.Dense(3)(x)
      clone = keras.models.Model(inputs, x)
      clone.compile(loss='mse', optimizer=keras.optimizers.RMSprop(lr=0.0001))
      clone.train_on_batch(input_arr, target_arr)

    keras_saved_model._assert_same_non_optimizer_objects(
        model, model_graph, clone, clone_graph)

  def testAssertModelCloneSameObjectsThrowError(self):
    input_arr = np.random.random((1, 3))
    target_arr = np.random.random((1, 3))

    model_graph = ops.Graph()
    clone_graph = ops.Graph()

    # Create two models with the same layers but different optimizers.
    with session.Session(graph=model_graph):
      inputs = keras.layers.Input(shape=(3,))
      x = keras.layers.Dense(2)(inputs)
      x = keras.layers.Dense(3)(x)
      model = keras.models.Model(inputs, x)

      model.compile(loss='mse', optimizer=training_module.AdadeltaOptimizer())
      model.train_on_batch(input_arr, target_arr)

    with session.Session(graph=clone_graph):
      inputs = keras.layers.Input(shape=(3,))
      x = keras.layers.Dense(2)(inputs)
      x = keras.layers.Dense(4)(x)
      x = keras.layers.Dense(3)(x)
      clone = keras.models.Model(inputs, x)
      clone.compile(loss='mse', optimizer=keras.optimizers.RMSprop(lr=0.0001))
      clone.train_on_batch(input_arr, target_arr)

  def testSaveSequentialModelWithoutInputShapes(self):
    model = sequential_model_without_input_shape(True)
    # A Sequential model that hasn't been built should raise an error.
    with self.assertRaisesRegexp(
        ValueError, 'Weights for sequential model have not yet been created'):
      keras_saved_model.export_saved_model(model, '')

    # Even with input_signature, the model's weights has not been created.
    with self.assertRaisesRegexp(
        ValueError, 'Weights for sequential model have not yet been created'):
      saved_model_dir = self._save_model_dir()
      keras_saved_model.export_saved_model(
          model,
          saved_model_dir,
          input_signature=tensor_spec.TensorSpec(
              shape=(10, 11, 12, 13, 14), dtype=dtypes.float32,
              name='spec_input'))

  @parameterized.parameters(
      {
          'model_builder': sequential_model_without_input_shape,
          'input_signature': [tensor_spec.TensorSpec(shape=[None, 3],
                                                     dtype=dtypes.float32)]},
      {
          'model_builder': subclassed_model,
          'input_signature': [tensor_spec.TensorSpec(shape=[None, 3],
                                                     dtype=dtypes.float32)]})
  def testServingOnly(self, model_builder, input_signature):
    if context.executing_eagerly():
      saved_model_dir = self._save_model_dir()
      input_arr = np.random.random((5, 3)).astype(np.float32)
      model = model_builder()
      ref_predict = model.predict(input_arr)

      keras_saved_model.export_saved_model(
          model,
          saved_model_dir,
          serving_only=True,
          input_signature=input_signature)

      # Load predict graph, and test predictions
      with session.Session(graph=ops.Graph()) as sess:
        inputs, outputs, _ = load_model(sess, saved_model_dir,
                                        mode_keys.ModeKeys.PREDICT)
        predictions = sess.run(outputs[next(iter(outputs.keys()))],
                               {inputs[next(iter(inputs.keys()))]: input_arr})
        self.assertAllClose(ref_predict, predictions, atol=1e-05)


if __name__ == '__main__':
  test.main()
