content: # Copyright 2017 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.
  # ==============================================================================
  """Tests for dnn.py with feature_column_v2."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import shutil
  import tempfile
  
  from unittest.mock import patch
  
  from absl.testing import parameterized
  import numpy as np
  import six
  import tensorflow as tf
  from tensorflow.core.example import example_pb2
  from tensorflow.core.example import feature_pb2
  from tensorflow.python.feature_column import feature_column_v2
  from tensorflow_estimator.python.estimator.canned import dnn
  from tensorflow_estimator.python.estimator.canned import dnn_testing_utils
  from tensorflow_estimator.python.estimator.canned import prediction_keys
  from tensorflow_estimator.python.estimator.export import export
  from tensorflow_estimator.python.estimator.inputs import numpy_io
  from tensorflow_estimator.python.estimator.inputs import pandas_io
  
  try:
    # pylint: disable=g-import-not-at-top
    import pandas as pd
    HAS_PANDAS = True
  except IOError:
    # Pandas writes a temporary file during import. If it fails, don't use pandas.
    HAS_PANDAS = False
  except ImportError:
    HAS_PANDAS = False
  
  
  def _dnn_classifier_fn(*args, **kwargs):
    return dnn.DNNClassifierV2(*args, **kwargs)
  
  
  class DNNModelFnV2Test(dnn_testing_utils.BaseDNNModelFnTest, tf.test.TestCase):
  
    def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
      tf.test.TestCase.__init__(self, methodName)
      dnn_testing_utils.BaseDNNModelFnTest.__init__(
          self, dnn.dnn_model_fn_v2, fc_impl=feature_column_v2)
  
  
  class DNNLogitFnV2Test(dnn_testing_utils.BaseDNNLogitFnTest, tf.test.TestCase):
  
    def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
      tf.test.TestCase.__init__(self, methodName)
      dnn_testing_utils.BaseDNNLogitFnTest.__init__(
          self, dnn.dnn_logit_fn_builder_v2, fc_impl=feature_column_v2)
  
  
  class DNNWarmStartingV2Test(dnn_testing_utils.BaseDNNWarmStartingTest,
                              tf.test.TestCase):
  
    def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
      tf.test.TestCase.__init__(self, methodName)
      dnn_testing_utils.BaseDNNWarmStartingTest.__init__(
          self, _dnn_classifier_fn, _dnn_regressor_fn, fc_impl=feature_column_v2)
  
  
  class DNNClassifierEvaluateV2Test(
      dnn_testing_utils.BaseDNNClassifierEvaluateTest, tf.test.TestCase):
  
    def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
      tf.test.TestCase.__init__(self, methodName)
      dnn_testing_utils.BaseDNNClassifierEvaluateTest.__init__(
          self, _dnn_classifier_fn, fc_impl=feature_column_v2)
  
  
  class DNNClassifierPredictV2Test(dnn_testing_utils.BaseDNNClassifierPredictTest,
                                   tf.test.TestCase):
  
    def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
      tf.test.TestCase.__init__(self, methodName)
      dnn_testing_utils.BaseDNNClassifierPredictTest.__init__(
          self, _dnn_classifier_fn, fc_impl=feature_column_v2)
  
  
  class DNNClassifierTrainV2Test(dnn_testing_utils.BaseDNNClassifierTrainTest,
                                 tf.test.TestCase):
  
    def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
      tf.test.TestCase.__init__(self, methodName)
      dnn_testing_utils.BaseDNNClassifierTrainTest.__init__(
          self, _dnn_classifier_fn, fc_impl=feature_column_v2)
  
  
  def _dnn_regressor_fn(*args, **kwargs):
    return dnn.DNNRegressorV2(*args, **kwargs)
  
  
  class DNNRegressorEvaluateV2Test(dnn_testing_utils.BaseDNNRegressorEvaluateTest,
                                   tf.test.TestCase):
  
    def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
      tf.test.TestCase.__init__(self, methodName)
      dnn_testing_utils.BaseDNNRegressorEvaluateTest.__init__(
          self, _dnn_regressor_fn, fc_impl=feature_column_v2)
  
  
  class DNNRegressorPredictV2Test(dnn_testing_utils.BaseDNNRegressorPredictTest,
                                  tf.test.TestCase):
  
    def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
      tf.test.TestCase.__init__(self, methodName)
      dnn_testing_utils.BaseDNNRegressorPredictTest.__init__(
          self, _dnn_regressor_fn, fc_impl=feature_column_v2)
  
  
  class DNNRegressorTrainV2Test(dnn_testing_utils.BaseDNNRegressorTrainTest,
                                tf.test.TestCase):
  
    def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
      tf.test.TestCase.__init__(self, methodName)
      dnn_testing_utils.BaseDNNRegressorTrainTest.__init__(
          self, _dnn_regressor_fn, fc_impl=feature_column_v2)
  
  
  def _queue_parsed_features(feature_map):
    tensors_to_enqueue = []
    keys = []
    for key, tensor in six.iteritems(feature_map):
      keys.append(key)
      tensors_to_enqueue.append(tensor)
    queue_dtypes = [x.dtype for x in tensors_to_enqueue]
    input_queue = tf.queue.FIFOQueue(capacity=100, dtypes=queue_dtypes)
    tf.compat.v1.train.queue_runner.add_queue_runner(
        tf.compat.v1.train.queue_runner.QueueRunner(
            input_queue, [input_queue.enqueue(tensors_to_enqueue)]))
    dequeued_tensors = input_queue.dequeue()
    return {keys[i]: dequeued_tensors[i] for i in range(len(dequeued_tensors))}
  
  
  class DNNRegressorIntegrationTest(tf.test.TestCase, parameterized.TestCase):
  
    def setUp(self):
      self._model_dir = tempfile.mkdtemp()
  
    def tearDown(self):
      if self._model_dir:
        tf.compat.v1.summary.FileWriterCache.clear()
        shutil.rmtree(self._model_dir)
  
    def _test_complete_flow(self, train_input_fn, eval_input_fn, predict_input_fn,
                            input_dimension, label_dimension, batch_size):
      feature_columns = [
          tf.feature_column.numeric_column('x', shape=(input_dimension,))
      ]
  
      est = dnn.DNNRegressorV2(
          hidden_units=(2, 2),
          feature_columns=feature_columns,
          label_dimension=label_dimension,
          model_dir=self._model_dir)
  
      # TRAIN
      num_steps = 10
      est.train(train_input_fn, steps=num_steps)
  
      # EVALUATE
      scores = est.evaluate(eval_input_fn)
      self.assertEqual(num_steps, scores[tf.compat.v1.GraphKeys.GLOBAL_STEP])
      self.assertIn('loss', six.iterkeys(scores))
  
      # PREDICT
      predictions = np.array([
          x[prediction_keys.PredictionKeys.PREDICTIONS]
          for x in est.predict(predict_input_fn)
      ])
      self.assertAllEqual((batch_size, label_dimension), predictions.shape)
  
      # EXPORT
      feature_spec = tf.feature_column.make_parse_example_spec(feature_columns)
      serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
          feature_spec)
      export_dir = est.export_saved_model(tempfile.mkdtemp(),
                                          serving_input_receiver_fn)
      self.assertTrue(tf.compat.v1.gfile.Exists(export_dir))
  
    def test_numpy_input_fn(self):
      """Tests complete flow with numpy_input_fn."""
      label_dimension = 2
      batch_size = 10
      data = np.linspace(0., 2., batch_size * label_dimension, dtype=np.float32)
      data = data.reshape(batch_size, label_dimension)
      # learn y = x
      train_input_fn = numpy_io.numpy_input_fn(
          x={'x': data},
          y=data,
          batch_size=batch_size,
          num_epochs=None,
          shuffle=True)
      eval_input_fn = numpy_io.numpy_input_fn(
          x={'x': data}, y=data, batch_size=batch_size, shuffle=False)
      predict_input_fn = numpy_io.numpy_input_fn(
          x={'x': data}, batch_size=batch_size, shuffle=False)
  
      self._test_complete_flow(
          train_input_fn=train_input_fn,
          eval_input_fn=eval_input_fn,
          predict_input_fn=predict_input_fn,
          input_dimension=label_dimension,
          label_dimension=label_dimension,
          batch_size=batch_size)
  
    def test_pandas_input_fn(self):
      """Tests complete flow with pandas_input_fn."""
      if not HAS_PANDAS:
        return
      label_dimension = 1
      batch_size = 10
      data = np.linspace(0., 2., batch_size, dtype=np.float32)
      x = pd.DataFrame({'x': data})
      y = pd.Series(data)
      train_input_fn = pandas_io.pandas_input_fn(
          x=x, y=y, batch_size=batch_size, num_epochs=None, shuffle=True)
      eval_input_fn = pandas_io.pandas_input_fn(
          x=x, y=y, batch_size=batch_size, shuffle=False)
      predict_input_fn = pandas_io.pandas_input_fn(
          x=x, batch_size=batch_size, shuffle=False)
  
      self._test_complete_flow(
          train_input_fn=train_input_fn,
          eval_input_fn=eval_input_fn,
          predict_input_fn=predict_input_fn,
          input_dimension=label_dimension,
          label_dimension=label_dimension,
          batch_size=batch_size)
  
    def test_input_fn_from_parse_example(self):
      """Tests complete flow with input_fn constructed from parse_example."""
      label_dimension = 2
      batch_size = 10
      data = np.linspace(0., 2., batch_size * label_dimension, dtype=np.float32)
      data = data.reshape(batch_size, label_dimension)
  
      serialized_examples = []
      for datum in data:
        example = example_pb2.Example(
            features=feature_pb2.Features(
                feature={
                    'x':
                        feature_pb2.Feature(
                            float_list=feature_pb2.FloatList(value=datum)),
                    'y':
                        feature_pb2.Feature(
                            float_list=feature_pb2.FloatList(value=datum)),
                }))
        serialized_examples.append(example.SerializeToString())
  
      feature_spec = {
          'x': tf.io.FixedLenFeature([label_dimension], tf.dtypes.float32),
          'y': tf.io.FixedLenFeature([label_dimension], tf.dtypes.float32),
      }
  
      def _train_input_fn():
        feature_map = tf.compat.v1.io.parse_example(serialized_examples,
                                                    feature_spec)
        features = _queue_parsed_features(feature_map)
        labels = features.pop('y')
        return features, labels
  
      def _eval_input_fn():
        feature_map = tf.compat.v1.io.parse_example(
            tf.compat.v1.train.limit_epochs(serialized_examples, num_epochs=1),
            feature_spec)
        features = _queue_parsed_features(feature_map)
        labels = features.pop('y')
        return features, labels
  
      def _predict_input_fn():
        feature_map = tf.compat.v1.io.parse_example(
            tf.compat.v1.train.limit_epochs(serialized_examples, num_epochs=1),
            feature_spec)
        features = _queue_parsed_features(feature_map)
        features.pop('y')
        return features, None
  
      self._test_complete_flow(
          train_input_fn=_train_input_fn,
          eval_input_fn=_eval_input_fn,
          predict_input_fn=_predict_input_fn,
          input_dimension=label_dimension,
          label_dimension=label_dimension,
          batch_size=batch_size)
  
  
  class DNNClassifierIntegrationTest(tf.test.TestCase):
  
    def setUp(self):
      self._model_dir = tempfile.mkdtemp()
  
    def tearDown(self):
      if self._model_dir:
        tf.compat.v1.summary.FileWriterCache.clear()
        shutil.rmtree(self._model_dir)
  
    def _as_label(self, data_in_float):
      return np.rint(data_in_float).astype(np.int64)
  
    def _test_complete_flow(self, train_input_fn, eval_input_fn, predict_input_fn,
                            input_dimension, n_classes, batch_size):
      feature_columns = [
          tf.feature_column.numeric_column('x', shape=(input_dimension,))
      ]
  
      est = dnn.DNNClassifierV2(
          hidden_units=(2, 2),
          feature_columns=feature_columns,
          n_classes=n_classes,
          model_dir=self._model_dir)
  
      # TRAIN
      num_steps = 10
      est.train(train_input_fn, steps=num_steps)
  
      # EVALUATE
      scores = est.evaluate(eval_input_fn)
      self.assertEqual(num_steps, scores[tf.compat.v1.GraphKeys.GLOBAL_STEP])
      self.assertIn('loss', six.iterkeys(scores))
  
      # PREDICT
      predicted_proba = np.array([
          x[prediction_keys.PredictionKeys.PROBABILITIES]
          for x in est.predict(predict_input_fn)
      ])
      self.assertAllEqual((batch_size, n_classes), predicted_proba.shape)
  
      # EXPORT
      feature_spec = tf.feature_column.make_parse_example_spec(feature_columns)
      serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
          feature_spec)
      export_dir = est.export_saved_model(tempfile.mkdtemp(),
                                          serving_input_receiver_fn)
      self.assertTrue(tf.compat.v1.gfile.Exists(export_dir))
  
    def test_numpy_input_fn(self):
      """Tests complete flow with numpy_input_fn."""
      n_classes = 3
      input_dimension = 2
      batch_size = 10
      data = np.linspace(
          0., n_classes - 1., batch_size * input_dimension, dtype=np.float32)
      x_data = data.reshape(batch_size, input_dimension)
      y_data = np.reshape(self._as_label(data[:batch_size]), (batch_size, 1))
      # learn y = x
      train_input_fn = numpy_io.numpy_input_fn(
          x={'x': x_data},
          y=y_data,
          batch_size=batch_size,
          num_epochs=None,
          shuffle=True)
      eval_input_fn = numpy_io.numpy_input_fn(
          x={'x': x_data}, y=y_data, batch_size=batch_size, shuffle=False)
      predict_input_fn = numpy_io.numpy_input_fn(
          x={'x': x_data}, batch_size=batch_size, shuffle=False)
  
      self._test_complete_flow(
          train_input_fn=train_input_fn,
          eval_input_fn=eval_input_fn,
          predict_input_fn=predict_input_fn,
          input_dimension=input_dimension,
          n_classes=n_classes,
          batch_size=batch_size)
  
    def test_pandas_input_fn(self):
      """Tests complete flow with pandas_input_fn."""
      if not HAS_PANDAS:
        return
      input_dimension = 1
      n_classes = 3
      batch_size = 10
      data = np.linspace(0., n_classes - 1., batch_size, dtype=np.float32)
      x = pd.DataFrame({'x': data})
      y = pd.Series(self._as_label(data))
      train_input_fn = pandas_io.pandas_input_fn(
          x=x, y=y, batch_size=batch_size, num_epochs=None, shuffle=True)
      eval_input_fn = pandas_io.pandas_input_fn(
          x=x, y=y, batch_size=batch_size, shuffle=False)
      predict_input_fn = pandas_io.pandas_input_fn(
          x=x, batch_size=batch_size, shuffle=False)
  
      self._test_complete_flow(
          train_input_fn=train_input_fn,
          eval_input_fn=eval_input_fn,
          predict_input_fn=predict_input_fn,
          input_dimension=input_dimension,
          n_classes=n_classes,
          batch_size=batch_size)
  
    def test_input_fn_from_parse_example(self):
      """Tests complete flow with input_fn constructed from parse_example."""
      input_dimension = 2
      n_classes = 3
      batch_size = 10
      data = np.linspace(
          0., n_classes - 1., batch_size * input_dimension, dtype=np.float32)
      data = data.reshape(batch_size, input_dimension)
  
      serialized_examples = []
      for datum in data:
        example = example_pb2.Example(
            features=feature_pb2.Features(
                feature={
                    'x':
                        feature_pb2.Feature(
                            float_list=feature_pb2.FloatList(value=datum)),
                    'y':
                        feature_pb2.Feature(
                            int64_list=feature_pb2.Int64List(
                                value=self._as_label(datum[:1]))),
                }))
        serialized_examples.append(example.SerializeToString())
  
      feature_spec = {
          'x': tf.io.FixedLenFeature([input_dimension], tf.dtypes.float32),
          'y': tf.io.FixedLenFeature([1], tf.dtypes.int64),
      }
  
      def _train_input_fn():
        feature_map = tf.compat.v1.io.parse_example(serialized_examples,
                                                    feature_spec)
        features = _queue_parsed_features(feature_map)
        labels = features.pop('y')
        return features, labels
  
      def _eval_input_fn():
        feature_map = tf.compat.v1.io.parse_example(
            tf.compat.v1.train.limit_epochs(serialized_examples, num_epochs=1),
            feature_spec)
        features = _queue_parsed_features(feature_map)
        labels = features.pop('y')
        return features, labels
  
      def _predict_input_fn():
        feature_map = tf.compat.v1.io.parse_example(
            tf.compat.v1.train.limit_epochs(serialized_examples, num_epochs=1),
            feature_spec)
        features = _queue_parsed_features(feature_map)
        features.pop('y')
        return features, None
  
      self._test_complete_flow(
          train_input_fn=_train_input_fn,
          eval_input_fn=_eval_input_fn,
          predict_input_fn=_predict_input_fn,
          input_dimension=input_dimension,
          n_classes=n_classes,
          batch_size=batch_size)
  
  
  class DNNTrainingMode(tf.test.TestCase):
    """Tests that training mode propagates to feature columns correctly."""
  
    def setUp(self):
      self._model_dir = tempfile.mkdtemp()
      self._label_dimension = 1
      self._batch_size = 10
  
    def tearDown(self):
      if self._model_dir:
        tf.compat.v1.summary.FileWriterCache.clear()
        shutil.rmtree(self._model_dir)
  
    def _create_data(self):
      data = np.linspace(
          0., 2., self._batch_size * self._label_dimension, dtype=np.float32)
      return data.reshape(self._batch_size, self._label_dimension)
  
    def _get_estimator(self):
      feature_columns = [
          tf.feature_column.numeric_column('x', shape=(self._label_dimension,))
      ]
      return dnn.DNNRegressorV2(
          hidden_units=(2, 2),
          feature_columns=feature_columns,
          label_dimension=self._label_dimension,
          model_dir=self._model_dir)
  
    def test_train_vs_eval_mode(self):
      data = self._create_data()
      train_input_fn = numpy_io.numpy_input_fn(
          x={'x': data},
          y=data,
          batch_size=self._batch_size,
          num_epochs=None,
          shuffle=True)
      eval_input_fn = numpy_io.numpy_input_fn(
          x={'x': data}, y=data, batch_size=self._batch_size, shuffle=False)
      est = self._get_estimator()
      with patch.object(
          tf.compat.v2.keras.layers.DenseFeatures, 'call',
          return_value=data) as mock_dense_features_call:
        est.train(train_input_fn, steps=10)
        est.evaluate(eval_input_fn)
      train_args, eval_args = mock_dense_features_call.call_args_list
      # DenseFeature should have been called with training = True in train.
      _, train_training_kwarg = train_args
      self.assertTrue(train_training_kwarg['training'])
      # DenseFeature should have been called with training = False in eval.
      _, eval_training_kwarg = eval_args
      self.assertFalse(eval_training_kwarg['training'])
  
  
  if __name__ == '__main__':
    tf.test.main()
  