Adds DNN-only tests for DNNLinearCombinedClassifier.

PiperOrigin-RevId: 158423119

tensorflow/python/estimator/BUILD

@@ -114,13 +114,10 @@ py_library(
     srcs = ["canned/dnn_testing_utils.py"],
     srcs_version = "PY2AND3",
     deps = [
-        ":dnn",
-        ":export_export",
         ":head",
         ":metric_keys",
         ":model_fn",
         ":numpy_io",
-        ":pandas_io",
         ":prediction_keys",
         "//tensorflow/core:protos_all_py",
         "//tensorflow/python:array_ops",
@@ -129,12 +126,9 @@ py_library(
         "//tensorflow/python:client_testlib",
         "//tensorflow/python:constant_op",
         "//tensorflow/python:control_flow_ops",
-        "//tensorflow/python:data_flow_ops",
         "//tensorflow/python:dtypes",
         "//tensorflow/python:framework_ops",
         "//tensorflow/python:math_ops",
-        "//tensorflow/python:parsing_ops",
-        "//tensorflow/python:platform",
         "//tensorflow/python:state_ops",
         "//tensorflow/python:summary",
         "//tensorflow/python:training",
@@ -153,28 +147,18 @@ py_test(
         ":dnn",
         ":dnn_testing_utils",
         ":export_export",
-        ":head",
         ":metric_keys",
-        ":model_fn",
         ":numpy_io",
         ":pandas_io",
         ":prediction_keys",
         "//tensorflow/core:protos_all_py",
-        "//tensorflow/python:array_ops",
-        "//tensorflow/python:check_ops",
-        "//tensorflow/python:client",
         "//tensorflow/python:client_testlib",
-        "//tensorflow/python:constant_op",
         "//tensorflow/python:data_flow_ops",
         "//tensorflow/python:dtypes",
         "//tensorflow/python:framework_ops",
-        "//tensorflow/python:math_ops",
         "//tensorflow/python:parsing_ops",
         "//tensorflow/python:platform",
-        "//tensorflow/python:state_ops",
-        "//tensorflow/python:summary",
         "//tensorflow/python:training",
-        "//tensorflow/python:variables",
         "//tensorflow/python/feature_column",
     ],
 )
@@ -188,10 +172,13 @@ py_library(
         ":head",
         ":model_fn",
         ":optimizers",
+        "//tensorflow/python:control_flow_ops",
+        "//tensorflow/python:framework_ops",
         "//tensorflow/python:init_ops",
         "//tensorflow/python:layers",
         "//tensorflow/python:nn",
         "//tensorflow/python:partitioned_variables",
+        "//tensorflow/python:state_ops",
         "//tensorflow/python:summary",
         "//tensorflow/python:training",
         "//tensorflow/python:variable_scope",
@@ -203,6 +190,7 @@ py_test(
     name = "dnn_linear_combined_test",
     size = "medium",
     srcs = ["canned/dnn_linear_combined_test.py"],
+    shard_count = 4,
     srcs_version = "PY2AND3",
     tags = ["no_pip"],
     deps = [

tensorflow/python/estimator/canned/dnn.py

@@ -39,8 +39,8 @@ _LEARNING_RATE = 0.05
 
 
 def _add_hidden_layer_summary(value, tag):
-  summary.scalar('%s_fraction_of_zero_values' % tag, nn.zero_fraction(value))
+  summary.scalar('%s/fraction_of_zero_values' % tag, nn.zero_fraction(value))
-  summary.histogram('%s_activation' % tag, value)
+  summary.histogram('%s/activation' % tag, value)
 
 
 def _dnn_model_fn(

tensorflow/python/estimator/canned/dnn_linear_combined.py

@@ -347,7 +347,8 @@ class DNNLinearCombinedClassifier(estimator.Estimator):
     """
     linear_feature_columns = linear_feature_columns or []
     dnn_feature_columns = dnn_feature_columns or []
-    self._feature_columns = linear_feature_columns + dnn_feature_columns
+    self._feature_columns = (
+        list(linear_feature_columns) + list(dnn_feature_columns))
     if not self._feature_columns:
       raise ValueError('Either linear_feature_columns or dnn_feature_columns '
                        'must be defined.')

tensorflow/python/estimator/canned/dnn_linear_combined_test.py

@@ -313,6 +313,54 @@ class DNNLinearCombinedRegressorIntegrationTest(test.TestCase):
         batch_size=batch_size)
 
 
+# A function to mimic dnn-classifier init reuse same tests.
+def _dnn_classifier_fn(
+    hidden_units,
+    feature_columns,
+    model_dir=None,
+    n_classes=2,
+    weight_feature_key=None,
+    optimizer='Adagrad',
+    config=None,
+    input_layer_partitioner=None):
+  return dnn_linear_combined.DNNLinearCombinedClassifier(
+      model_dir=model_dir,
+      dnn_hidden_units=hidden_units,
+      dnn_feature_columns=feature_columns,
+      dnn_optimizer=optimizer,
+      n_classes=n_classes,
+      weight_feature_key=weight_feature_key,
+      input_layer_partitioner=input_layer_partitioner,
+      config=config)
+
+
+class DNNOnlyClassifierEvaluateTest(
+    dnn_testing_utils.BaseDNNClassifierEvaluateTest, test.TestCase):
+
+  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
+    test.TestCase.__init__(self, methodName)
+    dnn_testing_utils.BaseDNNClassifierEvaluateTest.__init__(
+        self, _dnn_classifier_fn)
+
+
+class DNNOnlyClassifierPredictTest(
+    dnn_testing_utils.BaseDNNClassifierPredictTest, test.TestCase):
+
+  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
+    test.TestCase.__init__(self, methodName)
+    dnn_testing_utils.BaseDNNClassifierPredictTest.__init__(
+        self, _dnn_classifier_fn)
+
+
+class DNNOnlyClassifierTrainTest(
+    dnn_testing_utils.BaseDNNClassifierTrainTest, test.TestCase):
+
+  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
+    test.TestCase.__init__(self, methodName)
+    dnn_testing_utils.BaseDNNClassifierTrainTest.__init__(
+        self, _dnn_classifier_fn)
+
+
 class DNNLinearCombinedClassifierIntegrationTest(test.TestCase):
 
   def setUp(self):

tensorflow/python/estimator/canned/dnn_test.py

@@ -26,7 +26,6 @@ import six
 
 from tensorflow.core.example import example_pb2
 from tensorflow.core.example import feature_pb2
-from tensorflow.core.framework import summary_pb2
 from tensorflow.python.estimator.canned import dnn
 from tensorflow.python.estimator.canned import dnn_testing_utils
 from tensorflow.python.estimator.canned import metric_keys
@@ -41,11 +40,8 @@ from tensorflow.python.ops import data_flow_ops
 from tensorflow.python.ops import parsing_ops
 from tensorflow.python.platform import gfile
 from tensorflow.python.platform import test
-from tensorflow.python.summary import summary as summary_lib
-from tensorflow.python.training import checkpoint_utils
 from tensorflow.python.training import input as input_lib
 from tensorflow.python.training import queue_runner
-from tensorflow.python.training import session_run_hook
 
 try:
   # pylint: disable=g-import-not-at-top
@@ -58,6 +54,10 @@ except ImportError:
   HAS_PANDAS = False
 
 
+def _dnn_classifier_fn(*args, **kwargs):
+  return dnn.DNNClassifier(*args, **kwargs)
+
+
 class DNNModelFnTest(dnn_testing_utils.BaseDNNModelFnTest, test.TestCase):
 
   def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
@@ -132,82 +132,13 @@ class DNNRegressorEvaluateTest(test.TestCase):
     }, dnn_regressor.evaluate(input_fn=_input_fn, steps=1))
 
 
-class DNNClassifierEvaluateTest(test.TestCase):
+class DNNClassifierEvaluateTest(
+    dnn_testing_utils.BaseDNNClassifierEvaluateTest, test.TestCase):
 
-  def setUp(self):
+  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
-    self._model_dir = tempfile.mkdtemp()
+    test.TestCase.__init__(self, methodName)
-
+    dnn_testing_utils.BaseDNNClassifierEvaluateTest.__init__(
-  def tearDown(self):
+        self, _dnn_classifier_fn)
-    if self._model_dir:
-      shutil.rmtree(self._model_dir)
-
-  def test_one_dim(self):
-    """Asserts evaluation metrics for one-dimensional input and logits."""
-    global_step = 100
-    dnn_testing_utils.create_checkpoint(
-        (([[.6, .5]], [.1, -.1]), ([[1., .8], [-.8, -1.]], [.2, -.2]),
-         ([[-1.], [1.]], [.3]),), global_step, self._model_dir)
-
-    dnn_classifier = dnn.DNNClassifier(
-        hidden_units=(2, 2),
-        feature_columns=[feature_column.numeric_column('age')],
-        model_dir=self._model_dir)
-    def _input_fn():
-      # batch_size = 2, one false label, and one true.
-      return {'age': [[10.], [10.]]}, [[1], [0]]
-    # Uses identical numbers as DNNModelTest.test_one_dim_logits.
-    # See that test for calculation of logits.
-    # logits = [[-2.08], [-2.08]] =>
-    # logistic = 1/(1 + exp(-logits)) = [[0.11105597], [0.11105597]]
-    # loss = -1. * log(0.111) -1. * log(0.889) = 2.31544200
-    expected_loss = 2.31544200
-    self.assertAllClose({
-        metric_keys.MetricKeys.LOSS: expected_loss,
-        metric_keys.MetricKeys.LOSS_MEAN: expected_loss / 2.,
-        metric_keys.MetricKeys.ACCURACY: 0.5,
-        metric_keys.MetricKeys.PREDICTION_MEAN: 0.11105597,
-        metric_keys.MetricKeys.LABEL_MEAN: 0.5,
-        metric_keys.MetricKeys.ACCURACY_BASELINE: 0.5,
-        # There is no good way to calculate AUC for only two data points. But
-        # that is what the algorithm returns.
-        metric_keys.MetricKeys.AUC: 0.5,
-        metric_keys.MetricKeys.AUC_PR: 0.75,
-        ops.GraphKeys.GLOBAL_STEP: global_step
-    }, dnn_classifier.evaluate(input_fn=_input_fn, steps=1))
-
-  def test_multi_dim(self):
-    """Asserts evaluation metrics for multi-dimensional input and logits."""
-    global_step = 100
-    dnn_testing_utils.create_checkpoint(
-        (([[.6, .5], [-.6, -.5]], [.1, -.1]), ([[1., .8], [-.8, -1.]],
-                                               [.2, -.2]),
-         ([[-1., 1., .5], [-1., 1., .5]], [.3, -.3,
-                                           .0]),), global_step, self._model_dir)
-    n_classes = 3
-
-    dnn_classifier = dnn.DNNClassifier(
-        hidden_units=(2, 2),
-        feature_columns=[feature_column.numeric_column('age', shape=[2])],
-        n_classes=n_classes,
-        model_dir=self._model_dir)
-    def _input_fn():
-      # batch_size = 2, one false label, and one true.
-      return {'age': [[10., 8.], [10., 8.]]}, [[1], [0]]
-    # Uses identical numbers as
-    # DNNModelFnTest.test_multi_dim_input_multi_dim_logits.
-    # See that test for calculation of logits.
-    # logits = [[-0.48, 0.48, 0.39], [-0.48, 0.48, 0.39]]
-    # probabilities = exp(logits)/sum(exp(logits))
-    #               = [[0.16670536, 0.43538380, 0.39791084],
-    #                  [0.16670536, 0.43538380, 0.39791084]]
-    # loss = -log(0.43538380) - log(0.16670536)
-    expected_loss = 2.62305466
-    self.assertAllClose({
-        metric_keys.MetricKeys.LOSS: expected_loss,
-        metric_keys.MetricKeys.LOSS_MEAN: expected_loss / 2,
-        metric_keys.MetricKeys.ACCURACY: 0.5,
-        ops.GraphKeys.GLOBAL_STEP: global_step
-    }, dnn_classifier.evaluate(input_fn=_input_fn, steps=1))
 
 
 class DNNRegressorPredictTest(test.TestCase):
@@ -271,90 +202,13 @@ class DNNRegressorPredictTest(test.TestCase):
     }, next(dnn_regressor.predict(input_fn=input_fn)))
 
 
-class DNNClassifierPredictTest(test.TestCase):
+class DNNClassifierPredictTest(
+    dnn_testing_utils.BaseDNNClassifierPredictTest, test.TestCase):
 
-  def setUp(self):
+  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
-    self._model_dir = tempfile.mkdtemp()
+    test.TestCase.__init__(self, methodName)
-
+    dnn_testing_utils.BaseDNNClassifierPredictTest.__init__(
-  def tearDown(self):
+        self, _dnn_classifier_fn)
-    if self._model_dir:
-      shutil.rmtree(self._model_dir)
-
-  def test_one_dim(self):
-    """Asserts predictions for one-dimensional input and logits."""
-    dnn_testing_utils.create_checkpoint(
-        (([[.6, .5]], [.1, -.1]), ([[1., .8], [-.8, -1.]], [.2, -.2]),
-         ([[-1.], [1.]], [.3]),),
-        global_step=0,
-        model_dir=self._model_dir)
-
-    dnn_classifier = dnn.DNNClassifier(
-        hidden_units=(2, 2),
-        feature_columns=(feature_column.numeric_column('x'),),
-        model_dir=self._model_dir)
-    input_fn = numpy_io.numpy_input_fn(
-        x={'x': np.array([[10.]])}, batch_size=1, shuffle=False)
-    # Uses identical numbers as DNNModelTest.test_one_dim_logits.
-    # See that test for calculation of logits.
-    # logits = [-2.08] =>
-    # logistic = exp(-2.08)/(1 + exp(-2.08)) = 0.11105597
-    # probabilities = [1-logistic, logistic] = [0.88894403, 0.11105597]
-    # class_ids = argmax(probabilities) = [0]
-    predictions = next(dnn_classifier.predict(input_fn=input_fn))
-    self.assertAllClose([-2.08],
-                        predictions[prediction_keys.PredictionKeys.LOGITS])
-    self.assertAllClose([0.11105597],
-                        predictions[prediction_keys.PredictionKeys.LOGISTIC])
-    self.assertAllClose(
-        [0.88894403,
-         0.11105597], predictions[prediction_keys.PredictionKeys.PROBABILITIES])
-    self.assertAllClose([0],
-                        predictions[prediction_keys.PredictionKeys.CLASS_IDS])
-    self.assertAllEqual([b'0'],
-                        predictions[prediction_keys.PredictionKeys.CLASSES])
-
-  def test_multi_dim(self):
-    """Asserts predictions for multi-dimensional input and logits."""
-    dnn_testing_utils.create_checkpoint(
-        (([[.6, .5], [-.6, -.5]], [.1, -.1]),
-         ([[1., .8], [-.8, -1.]], [.2, -.2]), ([[-1., 1., .5], [-1., 1., .5]],
-                                               [.3, -.3, .0]),),
-        global_step=0,
-        model_dir=self._model_dir)
-
-    dnn_classifier = dnn.DNNClassifier(
-        hidden_units=(2, 2),
-        feature_columns=(feature_column.numeric_column('x', shape=(2,)),),
-        n_classes=3,
-        model_dir=self._model_dir)
-    input_fn = numpy_io.numpy_input_fn(
-        # Inputs shape is (batch_size, num_inputs).
-        x={'x': np.array([[10., 8.]])},
-        batch_size=1,
-        shuffle=False)
-    # Uses identical numbers as
-    # DNNModelFnTest.test_multi_dim_input_multi_dim_logits.
-    # See that test for calculation of logits.
-    # logits = [-0.48, 0.48, 0.39] =>
-    # probabilities[i] = exp(logits[i]) / sum_j exp(logits[j]) =>
-    # probabilities = [0.16670536, 0.43538380, 0.39791084]
-    # class_ids = argmax(probabilities) = [1]
-    predictions = next(dnn_classifier.predict(input_fn=input_fn))
-    self.assertItemsEqual(
-        [prediction_keys.PredictionKeys.LOGITS,
-         prediction_keys.PredictionKeys.PROBABILITIES,
-         prediction_keys.PredictionKeys.CLASS_IDS,
-         prediction_keys.PredictionKeys.CLASSES],
-        six.iterkeys(predictions))
-    self.assertAllClose(
-        [-0.48, 0.48, 0.39], predictions[prediction_keys.PredictionKeys.LOGITS])
-    self.assertAllClose(
-        [0.16670536, 0.43538380, 0.39791084],
-        predictions[prediction_keys.PredictionKeys.PROBABILITIES])
-    self.assertAllEqual(
-        [1], predictions[prediction_keys.PredictionKeys.CLASS_IDS])
-    self.assertAllEqual(
-        [b'1'], predictions[prediction_keys.PredictionKeys.CLASSES])
 
 
 def _queue_parsed_features(feature_map):
@@ -702,84 +556,6 @@ class DNNClassifierIntegrationTest(test.TestCase):
         batch_size=batch_size)
 
 
-class _SummaryHook(session_run_hook.SessionRunHook):
-  """Saves summaries every N steps."""
-
-  def __init__(self):
-    self._summaries = []
-
-  def begin(self):
-    self._summary_op = summary_lib.merge_all()
-
-  def before_run(self, run_context):
-    return session_run_hook.SessionRunArgs({'summary': self._summary_op})
-
-  def after_run(self, run_context, run_values):
-    s = summary_pb2.Summary()
-    s.ParseFromString(run_values.results['summary'])
-    self._summaries.append(s)
-
-  def summaries(self):
-    return tuple(self._summaries)
-
-
-def _assert_checkpoint(
-    testcase, global_step, input_units, hidden_units, output_units, model_dir):
-  """Asserts checkpoint contains expected variables with proper shapes.
-
-  Args:
-    testcase: A TestCase instance.
-    global_step: Expected global step value.
-    input_units: The dimension of input layer.
-    hidden_units: Iterable of integer sizes for the hidden layers.
-    output_units: The dimension of output layer (logits).
-    model_dir: The model directory.
-  """
-  shapes = {
-      name: shape
-      for (name, shape) in checkpoint_utils.list_variables(model_dir)
-  }
-
-  # Global step.
-  testcase.assertEqual([], shapes[ops.GraphKeys.GLOBAL_STEP])
-  testcase.assertEqual(
-      global_step,
-      checkpoint_utils.load_variable(
-          model_dir, ops.GraphKeys.GLOBAL_STEP))
-
-  # Hidden layer weights.
-  prev_layer_units = input_units
-  for i in range(len(hidden_units)):
-    layer_units = hidden_units[i]
-    testcase.assertAllEqual(
-        (prev_layer_units, layer_units),
-        shapes[dnn_testing_utils.HIDDEN_WEIGHTS_NAME_PATTERN % i])
-    testcase.assertAllEqual(
-        (layer_units,),
-        shapes[dnn_testing_utils.HIDDEN_BIASES_NAME_PATTERN % i])
-    prev_layer_units = layer_units
-
-  # Output layer weights.
-  testcase.assertAllEqual((prev_layer_units, output_units),
-                          shapes[dnn_testing_utils.LOGITS_WEIGHTS_NAME])
-  testcase.assertAllEqual((output_units,),
-                          shapes[dnn_testing_utils.LOGITS_BIASES_NAME])
-
-
-def _assert_simple_summary(testcase, expected_values, actual_summary):
-  """Assert summary the specified simple values.
-
-  Args:
-    testcase: A TestCase instance.
-    expected_values: Dict of expected tags and simple values.
-    actual_summary: `summary_pb2.Summary`.
-  """
-  testcase.assertAllClose(expected_values, {
-      v.tag: v.simple_value
-      for v in actual_summary.value if (v.tag in expected_values)
-  })
-
-
 class DNNRegressorTrainTest(test.TestCase):
 
   def setUp(self):
@@ -800,7 +576,7 @@ class DNNRegressorTrainTest(test.TestCase):
     num_steps = 5
     dnn_regressor.train(
         input_fn=lambda: ({'age': ((1,),)}, ((10,),)), steps=num_steps)
-    _assert_checkpoint(
+    dnn_testing_utils._assert_checkpoint(
         self, num_steps, input_units=1, hidden_units=hidden_units,
         output_units=1, model_dir=self._model_dir)
 
@@ -818,12 +594,12 @@ class DNNRegressorTrainTest(test.TestCase):
     # Train for a few steps, then validate optimizer, summaries, and
     # checkpoint.
     num_steps = 5
-    summary_hook = _SummaryHook()
+    summary_hook = dnn_testing_utils._SummaryHook()
     dnn_regressor.train(
         input_fn=lambda: ({'age': ((1,),)}, ((5.,),)), steps=num_steps,
         hooks=(summary_hook,))
     self.assertEqual(1, mock_optimizer.minimize.call_count)
-    _assert_checkpoint(
+    dnn_testing_utils._assert_checkpoint(
         self, num_steps, input_units=1, hidden_units=hidden_units,
         output_units=1, model_dir=self._model_dir)
     summaries = summary_hook.summaries()
@@ -858,7 +634,7 @@ class DNNRegressorTrainTest(test.TestCase):
     # Train for a few steps, then validate optimizer, summaries, and
     # checkpoint.
     num_steps = 5
-    summary_hook = _SummaryHook()
+    summary_hook = dnn_testing_utils._SummaryHook()
     dnn_regressor.train(
         input_fn=lambda: ({'age': [[10.]]}, [[1.]]), steps=num_steps,
         hooks=(summary_hook,))
@@ -866,17 +642,17 @@ class DNNRegressorTrainTest(test.TestCase):
     summaries = summary_hook.summaries()
     self.assertEqual(num_steps, len(summaries))
     for summary in summaries:
-      _assert_simple_summary(
+      dnn_testing_utils._assert_simple_summary(
           self,
           {
               metric_keys.MetricKeys.LOSS_MEAN: expected_loss,
-              'dnn/dnn/hiddenlayer_0_fraction_of_zero_values': 0.,
+              'dnn/dnn/hiddenlayer_0/fraction_of_zero_values': 0.,
-              'dnn/dnn/hiddenlayer_1_fraction_of_zero_values': 0.5,
+              'dnn/dnn/hiddenlayer_1/fraction_of_zero_values': 0.5,
-              'dnn/dnn/logits_fraction_of_zero_values': 0.,
+              'dnn/dnn/logits/fraction_of_zero_values': 0.,
               metric_keys.MetricKeys.LOSS: expected_loss,
           },
           summary)
-    _assert_checkpoint(
+    dnn_testing_utils._assert_checkpoint(
         self, base_global_step + num_steps, input_units=1,
         hidden_units=hidden_units, output_units=1, model_dir=self._model_dir)
 
@@ -912,7 +688,7 @@ class DNNRegressorTrainTest(test.TestCase):
     # Train for a few steps, then validate optimizer, summaries, and
     # checkpoint.
     num_steps = 5
-    summary_hook = _SummaryHook()
+    summary_hook = dnn_testing_utils._SummaryHook()
     dnn_regressor.train(
         input_fn=lambda: ({'age': [[10., 8.]]}, [[1., -1., 0.5]]),
         steps=num_steps,
@@ -921,187 +697,29 @@ class DNNRegressorTrainTest(test.TestCase):
     summaries = summary_hook.summaries()
     self.assertEqual(num_steps, len(summaries))
     for summary in summaries:
-      _assert_simple_summary(
+      dnn_testing_utils._assert_simple_summary(
           self,
           {
               metric_keys.MetricKeys.LOSS_MEAN: expected_loss / label_dimension,
-              'dnn/dnn/hiddenlayer_0_fraction_of_zero_values': 0.,
+              'dnn/dnn/hiddenlayer_0/fraction_of_zero_values': 0.,
-              'dnn/dnn/hiddenlayer_1_fraction_of_zero_values': 0.5,
+              'dnn/dnn/hiddenlayer_1/fraction_of_zero_values': 0.5,
-              'dnn/dnn/logits_fraction_of_zero_values': 0.,
+              'dnn/dnn/logits/fraction_of_zero_values': 0.,
               metric_keys.MetricKeys.LOSS: expected_loss,
           },
           summary)
-    _assert_checkpoint(
+    dnn_testing_utils._assert_checkpoint(
         self, base_global_step + num_steps, input_units=input_dimension,
         hidden_units=hidden_units, output_units=label_dimension,
         model_dir=self._model_dir)
 
 
-class DNNClassifierTrainTest(test.TestCase):
+class DNNClassifierTrainTest(
+    dnn_testing_utils.BaseDNNClassifierTrainTest, test.TestCase):
 
-  def setUp(self):
+  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
-    self._model_dir = tempfile.mkdtemp()
+    test.TestCase.__init__(self, methodName)
-
+    dnn_testing_utils.BaseDNNClassifierTrainTest.__init__(
-  def tearDown(self):
+        self, _dnn_classifier_fn)
-    if self._model_dir:
-      shutil.rmtree(self._model_dir)
-
-  def test_from_scratch_with_default_optimizer_binary(self):
-    hidden_units = (2, 2)
-    dnn_classifier = dnn.DNNClassifier(
-        hidden_units=hidden_units,
-        feature_columns=(feature_column.numeric_column('age'),),
-        model_dir=self._model_dir)
-
-    # Train for a few steps, then validate final checkpoint.
-    num_steps = 5
-    dnn_classifier.train(
-        input_fn=lambda: ({'age': [[10.]]}, [[1]]), steps=num_steps)
-    _assert_checkpoint(
-        self, num_steps, input_units=1, hidden_units=hidden_units,
-        output_units=1, model_dir=self._model_dir)
-
-  def test_from_scratch_with_default_optimizer_multi_class(self):
-    hidden_units = (2, 2)
-    n_classes = 3
-    dnn_classifier = dnn.DNNClassifier(
-        hidden_units=hidden_units,
-        feature_columns=(feature_column.numeric_column('age'),),
-        n_classes=n_classes,
-        model_dir=self._model_dir)
-
-    # Train for a few steps, then validate final checkpoint.
-    num_steps = 5
-    dnn_classifier.train(
-        input_fn=lambda: ({'age': [[10.]]}, [[2]]), steps=num_steps)
-    _assert_checkpoint(
-        self, num_steps, input_units=1, hidden_units=hidden_units,
-        output_units=n_classes, model_dir=self._model_dir)
-
-  def test_from_scratch_validate_summary(self):
-    hidden_units = (2, 2)
-    mock_optimizer = dnn_testing_utils.mock_optimizer(
-        self, hidden_units=hidden_units)
-    dnn_classifier = dnn.DNNClassifier(
-        hidden_units=hidden_units,
-        feature_columns=(feature_column.numeric_column('age'),),
-        optimizer=mock_optimizer,
-        model_dir=self._model_dir)
-    self.assertEqual(0, mock_optimizer.minimize.call_count)
-
-    # Train for a few steps, then validate optimizer, summaries, and
-    # checkpoint.
-    num_steps = 5
-    summary_hook = _SummaryHook()
-    dnn_classifier.train(
-        input_fn=lambda: ({'age': [[10.]]}, [[1]]), steps=num_steps,
-        hooks=(summary_hook,))
-    self.assertEqual(1, mock_optimizer.minimize.call_count)
-    _assert_checkpoint(
-        self, num_steps, input_units=1, hidden_units=hidden_units,
-        output_units=1, model_dir=self._model_dir)
-    summaries = summary_hook.summaries()
-    self.assertEqual(num_steps, len(summaries))
-    for summary in summaries:
-      summary_keys = [v.tag for v in summary.value]
-      self.assertIn(metric_keys.MetricKeys.LOSS, summary_keys)
-      self.assertIn(metric_keys.MetricKeys.LOSS_MEAN, summary_keys)
-
-  def test_binary_classification(self):
-    base_global_step = 100
-    hidden_units = (2, 2)
-    dnn_testing_utils.create_checkpoint(
-        (([[.6, .5]], [.1, -.1]), ([[1., .8], [-.8, -1.]], [.2, -.2]),
-         ([[-1.], [1.]], [.3]),), base_global_step, self._model_dir)
-
-    # Uses identical numbers as DNNModelFnTest.test_one_dim_logits.
-    # See that test for calculation of logits.
-    # logits = [-2.08] => probabilities = [0.889, 0.111]
-    # loss = -1. * log(0.111) = 2.19772100
-    expected_loss = 2.19772100
-    mock_optimizer = dnn_testing_utils.mock_optimizer(
-        self, hidden_units=hidden_units, expected_loss=expected_loss)
-    dnn_classifier = dnn.DNNClassifier(
-        hidden_units=hidden_units,
-        feature_columns=(feature_column.numeric_column('age'),),
-        optimizer=mock_optimizer,
-        model_dir=self._model_dir)
-    self.assertEqual(0, mock_optimizer.minimize.call_count)
-
-    # Train for a few steps, then validate optimizer, summaries, and
-    # checkpoint.
-    num_steps = 5
-    summary_hook = _SummaryHook()
-    dnn_classifier.train(
-        input_fn=lambda: ({'age': [[10.]]}, [[1]]), steps=num_steps,
-        hooks=(summary_hook,))
-    self.assertEqual(1, mock_optimizer.minimize.call_count)
-    summaries = summary_hook.summaries()
-    self.assertEqual(num_steps, len(summaries))
-    for summary in summaries:
-      _assert_simple_summary(
-          self,
-          {
-              metric_keys.MetricKeys.LOSS_MEAN: expected_loss,
-              'dnn/dnn/hiddenlayer_0_fraction_of_zero_values': 0.,
-              'dnn/dnn/hiddenlayer_1_fraction_of_zero_values': .5,
-              'dnn/dnn/logits_fraction_of_zero_values': 0.,
-              metric_keys.MetricKeys.LOSS: expected_loss,
-          },
-          summary)
-    _assert_checkpoint(
-        self, base_global_step + num_steps, input_units=1,
-        hidden_units=hidden_units, output_units=1, model_dir=self._model_dir)
-
-  def test_multi_class(self):
-    n_classes = 3
-    base_global_step = 100
-    hidden_units = (2, 2)
-    dnn_testing_utils.create_checkpoint(
-        (([[.6, .5]], [.1, -.1]), ([[1., .8], [-.8, -1.]], [.2, -.2]),
-         ([[-1., 1., .5], [-1., 1., .5]],
-          [.3, -.3, .0]),), base_global_step, self._model_dir)
-
-    # Uses identical numbers as DNNModelFnTest.test_multi_dim_logits.
-    # See that test for calculation of logits.
-    # logits = [-2.08, 2.08, 1.19] => probabilities = [0.0109, 0.7011, 0.2879]
-    # loss = -1. * log(0.7011) = 0.35505795
-    expected_loss = 0.35505795
-    mock_optimizer = dnn_testing_utils.mock_optimizer(
-        self, hidden_units=hidden_units, expected_loss=expected_loss)
-    dnn_classifier = dnn.DNNClassifier(
-        n_classes=n_classes,
-        hidden_units=hidden_units,
-        feature_columns=(feature_column.numeric_column('age'),),
-        optimizer=mock_optimizer,
-        model_dir=self._model_dir)
-    self.assertEqual(0, mock_optimizer.minimize.call_count)
-
-    # Train for a few steps, then validate optimizer, summaries, and
-    # checkpoint.
-    num_steps = 5
-    summary_hook = _SummaryHook()
-    dnn_classifier.train(
-        input_fn=lambda: ({'age': [[10.]]}, [[1]]), steps=num_steps,
-        hooks=(summary_hook,))
-    self.assertEqual(1, mock_optimizer.minimize.call_count)
-    summaries = summary_hook.summaries()
-    self.assertEqual(num_steps, len(summaries))
-    for summary in summaries:
-      _assert_simple_summary(
-          self,
-          {
-              metric_keys.MetricKeys.LOSS_MEAN: expected_loss,
-              'dnn/dnn/hiddenlayer_0_fraction_of_zero_values': 0.,
-              'dnn/dnn/hiddenlayer_1_fraction_of_zero_values': .5,
-              'dnn/dnn/logits_fraction_of_zero_values': 0.,
-              metric_keys.MetricKeys.LOSS: expected_loss,
-          },
-          summary)
-    _assert_checkpoint(
-        self, base_global_step + num_steps, input_units=1,
-        hidden_units=hidden_units, output_units=n_classes,
-        model_dir=self._model_dir)
 
 
 if __name__ == '__main__':

tensorflow/python/estimator/canned/dnn_testing_utils.py

@@ -25,9 +25,13 @@ import tempfile
 import numpy as np
 import six
 
+from tensorflow.core.framework import summary_pb2
 from tensorflow.python.client import session as tf_session
 from tensorflow.python.estimator import model_fn
 from tensorflow.python.estimator.canned import head as head_lib
+from tensorflow.python.estimator.canned import metric_keys
+from tensorflow.python.estimator.canned import prediction_keys
+from tensorflow.python.estimator.inputs import numpy_io
 from tensorflow.python.feature_column import feature_column
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
@@ -39,9 +43,12 @@ from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import state_ops
 from tensorflow.python.ops import variables as variables_lib
 from tensorflow.python.platform import test
+from tensorflow.python.summary import summary as summary_lib
+from tensorflow.python.training import checkpoint_utils
 from tensorflow.python.training import monitored_session
 from tensorflow.python.training import optimizer
 from tensorflow.python.training import saver
+from tensorflow.python.training import session_run_hook
 from tensorflow.python.training import training_util
 
 # pylint rules which are disabled by default for test files.
@@ -426,3 +433,421 @@ class BaseDNNModelFnTest(object):
             sess.run(estimator_spec.predictions)
           else:
             self.fail('Invalid mode: {}'.format(mode))
+
+
+class BaseDNNClassifierEvaluateTest(test.TestCase):
+
+  def __init__(self, dnn_classifier_fn):
+    self._dnn_classifier_fn = dnn_classifier_fn
+
+  def setUp(self):
+    self._model_dir = tempfile.mkdtemp()
+
+  def tearDown(self):
+    if self._model_dir:
+      shutil.rmtree(self._model_dir)
+
+  def test_one_dim(self):
+    """Asserts evaluation metrics for one-dimensional input and logits."""
+    global_step = 100
+    create_checkpoint(
+        (([[.6, .5]], [.1, -.1]), ([[1., .8], [-.8, -1.]], [.2, -.2]),
+         ([[-1.], [1.]], [.3]),), global_step, self._model_dir)
+
+    dnn_classifier = self._dnn_classifier_fn(
+        hidden_units=(2, 2),
+        feature_columns=[feature_column.numeric_column('age')],
+        model_dir=self._model_dir)
+    def _input_fn():
+      # batch_size = 2, one false label, and one true.
+      return {'age': [[10.], [10.]]}, [[1], [0]]
+    # Uses identical numbers as DNNModelTest.test_one_dim_logits.
+    # See that test for calculation of logits.
+    # logits = [[-2.08], [-2.08]] =>
+    # logistic = 1/(1 + exp(-logits)) = [[0.11105597], [0.11105597]]
+    # loss = -1. * log(0.111) -1. * log(0.889) = 2.31544200
+    expected_loss = 2.31544200
+    self.assertAllClose({
+        metric_keys.MetricKeys.LOSS: expected_loss,
+        metric_keys.MetricKeys.LOSS_MEAN: expected_loss / 2.,
+        metric_keys.MetricKeys.ACCURACY: 0.5,
+        metric_keys.MetricKeys.PREDICTION_MEAN: 0.11105597,
+        metric_keys.MetricKeys.LABEL_MEAN: 0.5,
+        metric_keys.MetricKeys.ACCURACY_BASELINE: 0.5,
+        # There is no good way to calculate AUC for only two data points. But
+        # that is what the algorithm returns.
+        metric_keys.MetricKeys.AUC: 0.5,
+        metric_keys.MetricKeys.AUC_PR: 0.75,
+        ops.GraphKeys.GLOBAL_STEP: global_step
+    }, dnn_classifier.evaluate(input_fn=_input_fn, steps=1))
+
+  def test_multi_dim(self):
+    """Asserts evaluation metrics for multi-dimensional input and logits."""
+    global_step = 100
+    create_checkpoint(
+        (([[.6, .5], [-.6, -.5]], [.1, -.1]), ([[1., .8], [-.8, -1.]],
+                                               [.2, -.2]),
+         ([[-1., 1., .5], [-1., 1., .5]], [.3, -.3,
+                                           .0]),), global_step, self._model_dir)
+    n_classes = 3
+
+    dnn_classifier = self._dnn_classifier_fn(
+        hidden_units=(2, 2),
+        feature_columns=[feature_column.numeric_column('age', shape=[2])],
+        n_classes=n_classes,
+        model_dir=self._model_dir)
+    def _input_fn():
+      # batch_size = 2, one false label, and one true.
+      return {'age': [[10., 8.], [10., 8.]]}, [[1], [0]]
+    # Uses identical numbers as
+    # DNNModelFnTest.test_multi_dim_input_multi_dim_logits.
+    # See that test for calculation of logits.
+    # logits = [[-0.48, 0.48, 0.39], [-0.48, 0.48, 0.39]]
+    # probabilities = exp(logits)/sum(exp(logits))
+    #               = [[0.16670536, 0.43538380, 0.39791084],
+    #                  [0.16670536, 0.43538380, 0.39791084]]
+    # loss = -log(0.43538380) - log(0.16670536)
+    expected_loss = 2.62305466
+    self.assertAllClose({
+        metric_keys.MetricKeys.LOSS: expected_loss,
+        metric_keys.MetricKeys.LOSS_MEAN: expected_loss / 2,
+        metric_keys.MetricKeys.ACCURACY: 0.5,
+        ops.GraphKeys.GLOBAL_STEP: global_step
+    }, dnn_classifier.evaluate(input_fn=_input_fn, steps=1))
+
+
+class BaseDNNClassifierPredictTest(test.TestCase):
+
+  def __init__(self, dnn_classifier_fn):
+    self._dnn_classifier_fn = dnn_classifier_fn
+
+  def setUp(self):
+    self._model_dir = tempfile.mkdtemp()
+
+  def tearDown(self):
+    if self._model_dir:
+      shutil.rmtree(self._model_dir)
+
+  def test_one_dim(self):
+    """Asserts predictions for one-dimensional input and logits."""
+    create_checkpoint(
+        (([[.6, .5]], [.1, -.1]), ([[1., .8], [-.8, -1.]], [.2, -.2]),
+         ([[-1.], [1.]], [.3]),),
+        global_step=0,
+        model_dir=self._model_dir)
+
+    dnn_classifier = self._dnn_classifier_fn(
+        hidden_units=(2, 2),
+        feature_columns=(feature_column.numeric_column('x'),),
+        model_dir=self._model_dir)
+    input_fn = numpy_io.numpy_input_fn(
+        x={'x': np.array([[10.]])}, batch_size=1, shuffle=False)
+    # Uses identical numbers as DNNModelTest.test_one_dim_logits.
+    # See that test for calculation of logits.
+    # logits = [-2.08] =>
+    # logistic = exp(-2.08)/(1 + exp(-2.08)) = 0.11105597
+    # probabilities = [1-logistic, logistic] = [0.88894403, 0.11105597]
+    # class_ids = argmax(probabilities) = [0]
+    predictions = next(dnn_classifier.predict(input_fn=input_fn))
+    self.assertAllClose([-2.08],
+                        predictions[prediction_keys.PredictionKeys.LOGITS])
+    self.assertAllClose([0.11105597],
+                        predictions[prediction_keys.PredictionKeys.LOGISTIC])
+    self.assertAllClose(
+        [0.88894403,
+         0.11105597], predictions[prediction_keys.PredictionKeys.PROBABILITIES])
+    self.assertAllClose([0],
+                        predictions[prediction_keys.PredictionKeys.CLASS_IDS])
+    self.assertAllEqual([b'0'],
+                        predictions[prediction_keys.PredictionKeys.CLASSES])
+
+  def test_multi_dim(self):
+    """Asserts predictions for multi-dimensional input and logits."""
+    create_checkpoint(
+        (([[.6, .5], [-.6, -.5]], [.1, -.1]),
+         ([[1., .8], [-.8, -1.]], [.2, -.2]), ([[-1., 1., .5], [-1., 1., .5]],
+                                               [.3, -.3, .0]),),
+        global_step=0,
+        model_dir=self._model_dir)
+
+    dnn_classifier = self._dnn_classifier_fn(
+        hidden_units=(2, 2),
+        feature_columns=(feature_column.numeric_column('x', shape=(2,)),),
+        n_classes=3,
+        model_dir=self._model_dir)
+    input_fn = numpy_io.numpy_input_fn(
+        # Inputs shape is (batch_size, num_inputs).
+        x={'x': np.array([[10., 8.]])},
+        batch_size=1,
+        shuffle=False)
+    # Uses identical numbers as
+    # DNNModelFnTest.test_multi_dim_input_multi_dim_logits.
+    # See that test for calculation of logits.
+    # logits = [-0.48, 0.48, 0.39] =>
+    # probabilities[i] = exp(logits[i]) / sum_j exp(logits[j]) =>
+    # probabilities = [0.16670536, 0.43538380, 0.39791084]
+    # class_ids = argmax(probabilities) = [1]
+    predictions = next(dnn_classifier.predict(input_fn=input_fn))
+    self.assertItemsEqual(
+        [prediction_keys.PredictionKeys.LOGITS,
+         prediction_keys.PredictionKeys.PROBABILITIES,
+         prediction_keys.PredictionKeys.CLASS_IDS,
+         prediction_keys.PredictionKeys.CLASSES],
+        six.iterkeys(predictions))
+    self.assertAllClose(
+        [-0.48, 0.48, 0.39], predictions[prediction_keys.PredictionKeys.LOGITS])
+    self.assertAllClose(
+        [0.16670536, 0.43538380, 0.39791084],
+        predictions[prediction_keys.PredictionKeys.PROBABILITIES])
+    self.assertAllEqual(
+        [1], predictions[prediction_keys.PredictionKeys.CLASS_IDS])
+    self.assertAllEqual(
+        [b'1'], predictions[prediction_keys.PredictionKeys.CLASSES])
+
+
+class _SummaryHook(session_run_hook.SessionRunHook):
+  """Saves summaries every N steps."""
+
+  def __init__(self):
+    self._summaries = []
+
+  def begin(self):
+    self._summary_op = summary_lib.merge_all()
+
+  def before_run(self, run_context):
+    return session_run_hook.SessionRunArgs({'summary': self._summary_op})
+
+  def after_run(self, run_context, run_values):
+    s = summary_pb2.Summary()
+    s.ParseFromString(run_values.results['summary'])
+    self._summaries.append(s)
+
+  def summaries(self):
+    return tuple(self._summaries)
+
+
+def _assert_checkpoint(
+    testcase, global_step, input_units, hidden_units, output_units, model_dir):
+  """Asserts checkpoint contains expected variables with proper shapes.
+
+  Args:
+    testcase: A TestCase instance.
+    global_step: Expected global step value.
+    input_units: The dimension of input layer.
+    hidden_units: Iterable of integer sizes for the hidden layers.
+    output_units: The dimension of output layer (logits).
+    model_dir: The model directory.
+  """
+  shapes = {
+      name: shape
+      for (name, shape) in checkpoint_utils.list_variables(model_dir)
+  }
+
+  # Global step.
+  testcase.assertEqual([], shapes[ops.GraphKeys.GLOBAL_STEP])
+  testcase.assertEqual(
+      global_step,
+      checkpoint_utils.load_variable(
+          model_dir, ops.GraphKeys.GLOBAL_STEP))
+
+  # Hidden layer weights.
+  prev_layer_units = input_units
+  for i in range(len(hidden_units)):
+    layer_units = hidden_units[i]
+    testcase.assertAllEqual(
+        (prev_layer_units, layer_units),
+        shapes[HIDDEN_WEIGHTS_NAME_PATTERN % i])
+    testcase.assertAllEqual(
+        (layer_units,),
+        shapes[HIDDEN_BIASES_NAME_PATTERN % i])
+    prev_layer_units = layer_units
+
+  # Output layer weights.
+  testcase.assertAllEqual((prev_layer_units, output_units),
+                          shapes[LOGITS_WEIGHTS_NAME])
+  testcase.assertAllEqual((output_units,),
+                          shapes[LOGITS_BIASES_NAME])
+
+
+def _assert_simple_summary(testcase, expected_values, actual_summary):
+  """Assert summary the specified simple values.
+
+  Args:
+    testcase: A TestCase instance.
+    expected_values: Dict of expected tags and simple values.
+    actual_summary: `summary_pb2.Summary`.
+  """
+  testcase.assertAllClose(expected_values, {
+      v.tag: v.simple_value
+      for v in actual_summary.value if (v.tag in expected_values)
+  })
+
+
+class BaseDNNClassifierTrainTest(test.TestCase):
+
+  def __init__(self, dnn_classifier_fn):
+    self._dnn_classifier_fn = dnn_classifier_fn
+
+  def setUp(self):
+    self._model_dir = tempfile.mkdtemp()
+
+  def tearDown(self):
+    if self._model_dir:
+      shutil.rmtree(self._model_dir)
+
+  def test_from_scratch_with_default_optimizer_binary(self):
+    hidden_units = (2, 2)
+    dnn_classifier = self._dnn_classifier_fn(
+        hidden_units=hidden_units,
+        feature_columns=(feature_column.numeric_column('age'),),
+        model_dir=self._model_dir)
+
+    # Train for a few steps, then validate final checkpoint.
+    num_steps = 5
+    dnn_classifier.train(
+        input_fn=lambda: ({'age': [[10.]]}, [[1]]), steps=num_steps)
+    _assert_checkpoint(
+        self, num_steps, input_units=1, hidden_units=hidden_units,
+        output_units=1, model_dir=self._model_dir)
+
+  def test_from_scratch_with_default_optimizer_multi_class(self):
+    hidden_units = (2, 2)
+    n_classes = 3
+    dnn_classifier = self._dnn_classifier_fn(
+        hidden_units=hidden_units,
+        feature_columns=(feature_column.numeric_column('age'),),
+        n_classes=n_classes,
+        model_dir=self._model_dir)
+
+    # Train for a few steps, then validate final checkpoint.
+    num_steps = 5
+    dnn_classifier.train(
+        input_fn=lambda: ({'age': [[10.]]}, [[2]]), steps=num_steps)
+    _assert_checkpoint(
+        self, num_steps, input_units=1, hidden_units=hidden_units,
+        output_units=n_classes, model_dir=self._model_dir)
+
+  def test_from_scratch_validate_summary(self):
+    hidden_units = (2, 2)
+    opt = mock_optimizer(
+        self, hidden_units=hidden_units)
+    dnn_classifier = self._dnn_classifier_fn(
+        hidden_units=hidden_units,
+        feature_columns=(feature_column.numeric_column('age'),),
+        optimizer=opt,
+        model_dir=self._model_dir)
+    self.assertEqual(0, opt.minimize.call_count)
+
+    # Train for a few steps, then validate optimizer, summaries, and
+    # checkpoint.
+    num_steps = 5
+    summary_hook = _SummaryHook()
+    dnn_classifier.train(
+        input_fn=lambda: ({'age': [[10.]]}, [[1]]), steps=num_steps,
+        hooks=(summary_hook,))
+    self.assertEqual(1, opt.minimize.call_count)
+    _assert_checkpoint(
+        self, num_steps, input_units=1, hidden_units=hidden_units,
+        output_units=1, model_dir=self._model_dir)
+    summaries = summary_hook.summaries()
+    self.assertEqual(num_steps, len(summaries))
+    for summary in summaries:
+      summary_keys = [v.tag for v in summary.value]
+      self.assertIn(metric_keys.MetricKeys.LOSS, summary_keys)
+      self.assertIn(metric_keys.MetricKeys.LOSS_MEAN, summary_keys)
+
+  def test_binary_classification(self):
+    base_global_step = 100
+    hidden_units = (2, 2)
+    create_checkpoint(
+        (([[.6, .5]], [.1, -.1]), ([[1., .8], [-.8, -1.]], [.2, -.2]),
+         ([[-1.], [1.]], [.3]),), base_global_step, self._model_dir)
+
+    # Uses identical numbers as DNNModelFnTest.test_one_dim_logits.
+    # See that test for calculation of logits.
+    # logits = [-2.08] => probabilities = [0.889, 0.111]
+    # loss = -1. * log(0.111) = 2.19772100
+    expected_loss = 2.19772100
+    opt = mock_optimizer(
+        self, hidden_units=hidden_units, expected_loss=expected_loss)
+    dnn_classifier = self._dnn_classifier_fn(
+        hidden_units=hidden_units,
+        feature_columns=(feature_column.numeric_column('age'),),
+        optimizer=opt,
+        model_dir=self._model_dir)
+    self.assertEqual(0, opt.minimize.call_count)
+
+    # Train for a few steps, then validate optimizer, summaries, and
+    # checkpoint.
+    num_steps = 5
+    summary_hook = _SummaryHook()
+    dnn_classifier.train(
+        input_fn=lambda: ({'age': [[10.]]}, [[1]]), steps=num_steps,
+        hooks=(summary_hook,))
+    self.assertEqual(1, opt.minimize.call_count)
+    summaries = summary_hook.summaries()
+    self.assertEqual(num_steps, len(summaries))
+    for summary in summaries:
+      _assert_simple_summary(
+          self,
+          {
+              metric_keys.MetricKeys.LOSS_MEAN: expected_loss,
+              'dnn/dnn/hiddenlayer_0/fraction_of_zero_values': 0.,
+              'dnn/dnn/hiddenlayer_1/fraction_of_zero_values': .5,
+              'dnn/dnn/logits/fraction_of_zero_values': 0.,
+              metric_keys.MetricKeys.LOSS: expected_loss,
+          },
+          summary)
+    _assert_checkpoint(
+        self, base_global_step + num_steps, input_units=1,
+        hidden_units=hidden_units, output_units=1, model_dir=self._model_dir)
+
+  def test_multi_class(self):
+    n_classes = 3
+    base_global_step = 100
+    hidden_units = (2, 2)
+    create_checkpoint(
+        (([[.6, .5]], [.1, -.1]), ([[1., .8], [-.8, -1.]], [.2, -.2]),
+         ([[-1., 1., .5], [-1., 1., .5]],
+          [.3, -.3, .0]),), base_global_step, self._model_dir)
+
+    # Uses identical numbers as DNNModelFnTest.test_multi_dim_logits.
+    # See that test for calculation of logits.
+    # logits = [-2.08, 2.08, 1.19] => probabilities = [0.0109, 0.7011, 0.2879]
+    # loss = -1. * log(0.7011) = 0.35505795
+    expected_loss = 0.35505795
+    opt = mock_optimizer(
+        self, hidden_units=hidden_units, expected_loss=expected_loss)
+    dnn_classifier = self._dnn_classifier_fn(
+        n_classes=n_classes,
+        hidden_units=hidden_units,
+        feature_columns=(feature_column.numeric_column('age'),),
+        optimizer=opt,
+        model_dir=self._model_dir)
+    self.assertEqual(0, opt.minimize.call_count)
+
+    # Train for a few steps, then validate optimizer, summaries, and
+    # checkpoint.
+    num_steps = 5
+    summary_hook = _SummaryHook()
+    dnn_classifier.train(
+        input_fn=lambda: ({'age': [[10.]]}, [[1]]), steps=num_steps,
+        hooks=(summary_hook,))
+    self.assertEqual(1, opt.minimize.call_count)
+    summaries = summary_hook.summaries()
+    self.assertEqual(num_steps, len(summaries))
+    for summary in summaries:
+      _assert_simple_summary(
+          self,
+          {
+              metric_keys.MetricKeys.LOSS_MEAN: expected_loss,
+              'dnn/dnn/hiddenlayer_0/fraction_of_zero_values': 0.,
+              'dnn/dnn/hiddenlayer_1/fraction_of_zero_values': .5,
+              'dnn/dnn/logits/fraction_of_zero_values': 0.,
+              metric_keys.MetricKeys.LOSS: expected_loss,
+          },
+          summary)
+    _assert_checkpoint(
+        self, base_global_step + num_steps, input_units=1,
+        hidden_units=hidden_units, output_units=n_classes,
+        model_dir=self._model_dir)