Fix some bugs in saved_model_export_utils for classification.

Change: 145704327

2 files changed, 185 insertions, 13 deletions

diff --git a/tensorflow/contrib/learn/python/learn/utils/saved_model_export_utils.py b/tensorflow/contrib/learn/python/learn/utils/saved_model_export_utils.py
index 296745e5c9..53147a39e1 100644
--- a/tensorflow/contrib/learn/python/learn/utils/saved_model_export_utils.py
+++ b/tensorflow/contrib/learn/python/learn/utils/saved_model_export_utils.py
@@ -85,8 +85,8 @@ def build_standardized_signature_def(
   if _is_classification_problem(problem_type, input_tensors, output_tensors):
     (_, examples), = input_tensors.items()
     classes = output_tensors.get(prediction_key.PredictionKey.CLASSES)
-    scores = output_tensors.get(prediction_key.PredictionKey.SCORES)
-    if not (classes or scores):
+    scores = _get_classification_scores(output_tensors)
+    if classes is None and scores is None:
       (_, classes), = output_tensors.items()
     return signature_def_utils.classification_signature_def(
         examples, classes, scores)
@@ -99,13 +99,22 @@ def build_standardized_signature_def(
         input_tensors, output_tensors)
 
 
+def _get_classification_scores(output_tensors):
+  scores = output_tensors.get(prediction_key.PredictionKey.SCORES)
+  if scores is None:
+    scores = output_tensors.get(prediction_key.PredictionKey.PROBABILITIES)
+  return scores
+
+
 def _is_classification_problem(problem_type, input_tensors, output_tensors):
   classes = output_tensors.get(prediction_key.PredictionKey.CLASSES)
-  scores = output_tensors.get(prediction_key.PredictionKey.SCORES)
+  scores = _get_classification_scores(output_tensors)
   return ((problem_type == constants.ProblemType.CLASSIFICATION or
            problem_type == constants.ProblemType.LOGISTIC_REGRESSION)
           and len(input_tensors) == 1
-          and (classes or scores or len(output_tensors) == 1))
+          and (classes is not None or
+               scores is not None or
+               len(output_tensors) == 1))
 
 
 def _is_regression_problem(problem_type, input_tensors, output_tensors):
diff --git a/tensorflow/contrib/learn/python/learn/utils/saved_model_export_utils_test.py b/tensorflow/contrib/learn/python/learn/utils/saved_model_export_utils_test.py
index 5ff2cf76f9..45f95a0163 100644
--- a/tensorflow/contrib/learn/python/learn/utils/saved_model_export_utils_test.py
+++ b/tensorflow/contrib/learn/python/learn/utils/saved_model_export_utils_test.py
@@ -49,7 +49,7 @@ from tensorflow.python.util import compat
 
 class SavedModelExportUtilsTest(test.TestCase):
 
-  def test_build_standardized_signature_def(self):
+  def test_build_standardized_signature_def_regression(self):
     input_tensors = {
         "input-1":
             array_ops.placeholder(
@@ -61,26 +61,189 @@ class SavedModelExportUtilsTest(test.TestCase):
                 dtypes.float32, 1, name="output-tensor-1")
     }
     problem_type = constants.ProblemType.LINEAR_REGRESSION
-    regression_signature_def = (
+    actual_signature_def = (
         saved_model_export_utils.build_standardized_signature_def(
             input_tensors, output_tensors, problem_type))
-    expected_regression_signature_def = meta_graph_pb2.SignatureDef()
+    expected_signature_def = meta_graph_pb2.SignatureDef()
     shape = tensor_shape_pb2.TensorShapeProto(
         dim=[tensor_shape_pb2.TensorShapeProto.Dim(size=1)])
     dtype = types_pb2.DataType.Value("DT_FLOAT")
-    expected_regression_signature_def.inputs[
+    expected_signature_def.inputs[
         signature_constants.REGRESS_INPUTS].CopyFrom(
             meta_graph_pb2.TensorInfo(
                 name="input-tensor-1:0", dtype=dtype, tensor_shape=shape))
-    expected_regression_signature_def.outputs[
+    expected_signature_def.outputs[
         signature_constants.REGRESS_OUTPUTS].CopyFrom(
             meta_graph_pb2.TensorInfo(
                 name="output-tensor-1:0", dtype=dtype, tensor_shape=shape))
 
-    expected_regression_signature_def.method_name = (
-        signature_constants.REGRESS_METHOD_NAME)
-    self.assertEqual(regression_signature_def,
-                     expected_regression_signature_def)
+    expected_signature_def.method_name = signature_constants.REGRESS_METHOD_NAME
+    self.assertEqual(actual_signature_def, expected_signature_def)
+
+  def test_build_standardized_signature_def_classification(self):
+    """Tests classification with one output tensor."""
+    input_tensors = {
+        "input-1":
+            array_ops.placeholder(
+                dtypes.float32, 1, name="input-tensor-1")
+    }
+    output_tensors = {
+        "output-1":
+            array_ops.placeholder(
+                dtypes.string, 1, name="output-tensor-1")
+    }
+    problem_type = constants.ProblemType.CLASSIFICATION
+    actual_signature_def = (
+        saved_model_export_utils.build_standardized_signature_def(
+            input_tensors, output_tensors, problem_type))
+    expected_signature_def = meta_graph_pb2.SignatureDef()
+    shape = tensor_shape_pb2.TensorShapeProto(
+        dim=[tensor_shape_pb2.TensorShapeProto.Dim(size=1)])
+    dtype_float = types_pb2.DataType.Value("DT_FLOAT")
+    dtype_string = types_pb2.DataType.Value("DT_STRING")
+    expected_signature_def.inputs[
+        signature_constants.CLASSIFY_INPUTS].CopyFrom(
+            meta_graph_pb2.TensorInfo(
+                name="input-tensor-1:0", dtype=dtype_float, tensor_shape=shape))
+    expected_signature_def.outputs[
+        signature_constants.CLASSIFY_OUTPUT_CLASSES].CopyFrom(
+            meta_graph_pb2.TensorInfo(
+                name="output-tensor-1:0", dtype=dtype_string,
+                tensor_shape=shape))
+
+    expected_signature_def.method_name = (
+        signature_constants.CLASSIFY_METHOD_NAME)
+    self.assertEqual(actual_signature_def, expected_signature_def)
+
+  def test_build_standardized_signature_def_classification2(self):
+    """Tests multiple output tensors that include classes and probabilites."""
+    input_tensors = {
+        "input-1":
+            array_ops.placeholder(
+                dtypes.float32, 1, name="input-tensor-1")
+    }
+    output_tensors = {
+        "classes":
+            array_ops.placeholder(
+                dtypes.string, 1, name="output-tensor-classes"),
+        # Will be used for CLASSIFY_OUTPUT_SCORES.
+        "probabilities":
+            array_ops.placeholder(
+                dtypes.float32, 1, name="output-tensor-proba"),
+        "logits":
+            array_ops.placeholder(
+                dtypes.float32, 1, name="output-tensor-logits-unused"),
+    }
+    problem_type = constants.ProblemType.CLASSIFICATION
+    actual_signature_def = (
+        saved_model_export_utils.build_standardized_signature_def(
+            input_tensors, output_tensors, problem_type))
+    expected_signature_def = meta_graph_pb2.SignatureDef()
+    shape = tensor_shape_pb2.TensorShapeProto(
+        dim=[tensor_shape_pb2.TensorShapeProto.Dim(size=1)])
+    dtype_float = types_pb2.DataType.Value("DT_FLOAT")
+    dtype_string = types_pb2.DataType.Value("DT_STRING")
+    expected_signature_def.inputs[
+        signature_constants.CLASSIFY_INPUTS].CopyFrom(
+            meta_graph_pb2.TensorInfo(
+                name="input-tensor-1:0", dtype=dtype_float, tensor_shape=shape))
+    expected_signature_def.outputs[
+        signature_constants.CLASSIFY_OUTPUT_CLASSES].CopyFrom(
+            meta_graph_pb2.TensorInfo(
+                name="output-tensor-classes:0", dtype=dtype_string,
+                tensor_shape=shape))
+    expected_signature_def.outputs[
+        signature_constants.CLASSIFY_OUTPUT_SCORES].CopyFrom(
+            meta_graph_pb2.TensorInfo(
+                name="output-tensor-proba:0", dtype=dtype_float,
+                tensor_shape=shape))
+
+    expected_signature_def.method_name = (
+        signature_constants.CLASSIFY_METHOD_NAME)
+    self.assertEqual(actual_signature_def, expected_signature_def)
+
+  def test_build_standardized_signature_def_classification3(self):
+    """Tests multiple output tensors that include classes and scores."""
+    input_tensors = {
+        "input-1":
+            array_ops.placeholder(
+                dtypes.float32, 1, name="input-tensor-1")
+    }
+    output_tensors = {
+        "classes":
+            array_ops.placeholder(
+                dtypes.string, 1, name="output-tensor-classes"),
+        "scores":
+            array_ops.placeholder(
+                dtypes.float32, 1, name="output-tensor-scores"),
+        "logits":
+            array_ops.placeholder(
+                dtypes.float32, 1, name="output-tensor-logits-unused"),
+    }
+    problem_type = constants.ProblemType.CLASSIFICATION
+    actual_signature_def = (
+        saved_model_export_utils.build_standardized_signature_def(
+            input_tensors, output_tensors, problem_type))
+    expected_signature_def = meta_graph_pb2.SignatureDef()
+    shape = tensor_shape_pb2.TensorShapeProto(
+        dim=[tensor_shape_pb2.TensorShapeProto.Dim(size=1)])
+    dtype_float = types_pb2.DataType.Value("DT_FLOAT")
+    dtype_string = types_pb2.DataType.Value("DT_STRING")
+    expected_signature_def.inputs[
+        signature_constants.CLASSIFY_INPUTS].CopyFrom(
+            meta_graph_pb2.TensorInfo(
+                name="input-tensor-1:0", dtype=dtype_float, tensor_shape=shape))
+    expected_signature_def.outputs[
+        signature_constants.CLASSIFY_OUTPUT_CLASSES].CopyFrom(
+            meta_graph_pb2.TensorInfo(
+                name="output-tensor-classes:0", dtype=dtype_string,
+                tensor_shape=shape))
+    expected_signature_def.outputs[
+        signature_constants.CLASSIFY_OUTPUT_SCORES].CopyFrom(
+            meta_graph_pb2.TensorInfo(
+                name="output-tensor-scores:0", dtype=dtype_float,
+                tensor_shape=shape))
+
+    expected_signature_def.method_name = (
+        signature_constants.CLASSIFY_METHOD_NAME)
+    self.assertEqual(actual_signature_def, expected_signature_def)
+
+  def test_build_standardized_signature_def_classification4(self):
+    """Tests classification without classes tensor."""
+    input_tensors = {
+        "input-1":
+            array_ops.placeholder(
+                dtypes.float32, 1, name="input-tensor-1")
+    }
+    output_tensors = {
+        "probabilities":
+            array_ops.placeholder(
+                dtypes.float32, 1, name="output-tensor-proba"),
+        "logits":
+            array_ops.placeholder(
+                dtypes.float32, 1, name="output-tensor-logits-unused"),
+    }
+    problem_type = constants.ProblemType.CLASSIFICATION
+    actual_signature_def = (
+        saved_model_export_utils.build_standardized_signature_def(
+            input_tensors, output_tensors, problem_type))
+    expected_signature_def = meta_graph_pb2.SignatureDef()
+    shape = tensor_shape_pb2.TensorShapeProto(
+        dim=[tensor_shape_pb2.TensorShapeProto.Dim(size=1)])
+    dtype_float = types_pb2.DataType.Value("DT_FLOAT")
+    expected_signature_def.inputs[
+        signature_constants.CLASSIFY_INPUTS].CopyFrom(
+            meta_graph_pb2.TensorInfo(
+                name="input-tensor-1:0", dtype=dtype_float, tensor_shape=shape))
+    expected_signature_def.outputs[
+        signature_constants.CLASSIFY_OUTPUT_SCORES].CopyFrom(
+            meta_graph_pb2.TensorInfo(
+                name="output-tensor-proba:0", dtype=dtype_float,
+                tensor_shape=shape))
+
+    expected_signature_def.method_name = (
+        signature_constants.CLASSIFY_METHOD_NAME)
+    self.assertEqual(actual_signature_def, expected_signature_def)
 
   def test_get_input_alternatives(self):
     input_ops = input_fn_utils.InputFnOps("bogus features dict", None,