Move functions that depend on problem type (loss, metric, predictions) from base class to Classifier and Regressor.

Change: 123818693

1 files changed, 126 insertions, 50 deletions

diff --git a/tensorflow/contrib/learn/python/learn/estimators/dnn_linear_combined.py b/tensorflow/contrib/learn/python/learn/estimators/dnn_linear_combined.py
index ffbe052355..8acf2c5738 100644
--- a/tensorflow/contrib/learn/python/learn/estimators/dnn_linear_combined.py
+++ b/tensorflow/contrib/learn/python/learn/estimators/dnn_linear_combined.py
@@ -85,7 +85,6 @@ class _DNNLinearCombinedBaseEstimator(estimator.BaseEstimator):
 
   def __init__(self,
                model_dir=None,
-               n_classes=2,
                weight_column_name=None,
                linear_feature_columns=None,
                linear_optimizer=None,
@@ -96,7 +95,6 @@ class _DNNLinearCombinedBaseEstimator(estimator.BaseEstimator):
                config=None):
     super(_DNNLinearCombinedBaseEstimator, self).__init__(model_dir=model_dir,
                                                           config=config)
-    self._n_classes = n_classes
     self._weight_column_name = weight_column_name
     self._linear_feature_columns = linear_feature_columns
     self._linear_optimizer = linear_optimizer
@@ -121,12 +119,7 @@ class _DNNLinearCombinedBaseEstimator(estimator.BaseEstimator):
     Returns:
       Numpy array of predicted classes or regression values.
     """
-    predictions = self._infer_model(x=x,
-                                    input_fn=input_fn,
-                                    batch_size=batch_size)
-    if self._n_classes > 1:
-      predictions = np.argmax(predictions, axis=1)
-    return predictions
+    raise NotImplementedError
 
   def predict_proba(self, x=None, input_fn=None, batch_size=None):
     """Returns prediction probabilities for given features (classification).
@@ -139,7 +132,7 @@ class _DNNLinearCombinedBaseEstimator(estimator.BaseEstimator):
     Returns:
       Numpy array of predicted probabilities.
     """
-    return self._infer_model(x=x, input_fn=input_fn, batch_size=batch_size)
+    raise NotImplementedError
 
   @property
   def linear_weights_(self):
@@ -196,26 +189,7 @@ class _DNNLinearCombinedBaseEstimator(estimator.BaseEstimator):
     return result
 
   def _get_eval_ops(self, features, targets, metrics=None):
-    """See base class."""
-    logits = self._logits(features)
-    result = {"loss": metrics_lib.streaming_mean(self._loss(
-        logits, targets,
-        weight_tensor=self._get_weight_tensor(features)))}
-
-    # Adding default metrics
-    if metrics is None and self._n_classes > 1:
-      metrics = {"accuracy": metrics_lib.streaming_accuracy}
-
-    if self._n_classes == 2:
-      predictions = math_ops.sigmoid(logits)
-      result["auc"] = metrics_lib.streaming_auc(predictions, targets)
-
-    if metrics:
-      predictions = self._logits_to_predictions(logits, proba=False)
-      result.update(self._run_metrics(predictions, targets, metrics,
-                                      self._get_weight_tensor(features)))
-
-    return result
+    raise NotImplementedError
 
   def _get_predict_ops(self, features):
     """See base class."""
@@ -223,16 +197,7 @@ class _DNNLinearCombinedBaseEstimator(estimator.BaseEstimator):
     return self._logits_to_predictions(logits, proba=True)
 
   def _logits_to_predictions(self, logits, proba=False):
-    if self._n_classes < 2:
-      return array_ops.reshape(logits, [-1])
-
-    if self._n_classes == 2:
-      logits = array_ops.concat(1, [array_ops.zeros_like(logits), logits])
-
-    if proba:
-      return nn.softmax(logits)
-    else:
-      return math_ops.argmax(logits, 1)
+    raise NotImplementedError
 
   def _get_feature_ops_from_example(self, examples_batch):
     column_types = layers.create_feature_spec_for_parsing((
@@ -242,7 +207,7 @@ class _DNNLinearCombinedBaseEstimator(estimator.BaseEstimator):
     return features
 
   def _num_label_columns(self):
-    return 1 if self._n_classes <= 2 else self._n_classes
+    raise NotImplementedError
 
   def _get_linear_feature_columns(self):
     return sorted(
@@ -340,15 +305,11 @@ class _DNNLinearCombinedBaseEstimator(estimator.BaseEstimator):
           math_ops.to_float(features[self._weight_column_name]),
           shape=(-1,))
 
+  def _loss_vec(self, logits, targets):
+    raise NotImplementedError
+
   def _loss(self, logits, target, weight_tensor):
-    if self._n_classes < 2:
-      loss_vec = math_ops.square(logits - math_ops.to_float(target))
-    elif self._n_classes == 2:
-      loss_vec = nn.sigmoid_cross_entropy_with_logits(
-          logits, array_ops.reshape(math_ops.to_float(target), [-1, 1]))
-    else:
-      loss_vec = nn.sparse_softmax_cross_entropy_with_logits(
-          logits, array_ops.reshape(target, [-1]))
+    loss_vec = self._loss_vec(logits, target)
 
     if weight_tensor is None:
       return math_ops.reduce_mean(loss_vec, name="loss")
@@ -486,13 +447,14 @@ class DNNLinearCombinedClassifier(_DNNLinearCombinedBaseEstimator):
                dnn_hidden_units=None,
                dnn_activation_fn=nn.relu,
                config=None):
+
     if n_classes < 2:
       raise ValueError("n_classes should be greater than 1. Given: {}".format(
           n_classes))
 
+    self._n_classes = n_classes
     super(DNNLinearCombinedClassifier, self).__init__(
         model_dir=model_dir,
-        n_classes=n_classes,
         weight_column_name=weight_column_name,
         linear_feature_columns=linear_feature_columns,
         linear_optimizer=linear_optimizer,
@@ -502,6 +464,80 @@ class DNNLinearCombinedClassifier(_DNNLinearCombinedBaseEstimator):
         dnn_activation_fn=dnn_activation_fn,
         config=config)
 
+  def predict(self, x=None, input_fn=None, batch_size=None):
+    """Returns predictions for given features.
+
+    Args:
+      x: features.
+      input_fn: Input function. If set, x must be None.
+      batch_size: Override default batch size.
+
+    Returns:
+      Numpy array of predicted classes or regression values.
+    """
+    predictions = self._infer_model(x=x,
+                                    input_fn=input_fn,
+                                    batch_size=batch_size)
+    predictions = np.argmax(predictions, axis=1)
+    return predictions
+
+  def predict_proba(self, x=None, input_fn=None, batch_size=None):
+    """Returns prediction probabilities for given features.
+
+    Args:
+      x: features.
+      input_fn: Input function. If set, x and y must be None.
+      batch_size: Override default batch size.
+
+    Returns:
+      Numpy array of predicted probabilities.
+    """
+    return self._infer_model(x=x, input_fn=input_fn, batch_size=batch_size)
+
+  def _loss_vec(self, logits, target):
+    if self._n_classes == 2:
+      loss_vec = nn.sigmoid_cross_entropy_with_logits(
+          logits, array_ops.reshape(math_ops.to_float(target), [-1, 1]))
+    else:
+      loss_vec = nn.sparse_softmax_cross_entropy_with_logits(
+          logits, array_ops.reshape(target, [-1]))
+
+    return loss_vec
+
+  def _logits_to_predictions(self, logits, proba=False):
+    if self._n_classes == 2:
+      logits = array_ops.concat(1, [array_ops.zeros_like(logits), logits])
+
+    if proba:
+      return nn.softmax(logits)
+    else:
+      return math_ops.argmax(logits, 1)
+
+  def _num_label_columns(self):
+    return 1 if self._n_classes == 2 else self._n_classes
+
+  def _get_eval_ops(self, features, targets, metrics=None):
+    """See base class."""
+    logits = self._logits(features)
+    result = {"loss": metrics_lib.streaming_mean(self._loss(
+        logits, targets,
+        weight_tensor=self._get_weight_tensor(features)))}
+
+    # Adding default metrics
+    if metrics is None:
+      metrics = {"accuracy": metrics_lib.streaming_accuracy}
+
+    if self._n_classes == 2:
+      predictions = math_ops.sigmoid(logits)
+      result["eval_auc"] = metrics_lib.streaming_auc(predictions, targets)
+
+    if metrics:
+      predictions = self._logits_to_predictions(logits, proba=False)
+      result.update(self._run_metrics(predictions, targets, metrics,
+                                      self._get_weight_tensor(features)))
+
+    return result
+
 
 class DNNLinearCombinedRegressor(_DNNLinearCombinedBaseEstimator):
   """A regressor for TensorFlow Linear and DNN joined training models.
@@ -588,7 +624,6 @@ class DNNLinearCombinedRegressor(_DNNLinearCombinedBaseEstimator):
                config=None):
     super(DNNLinearCombinedRegressor, self).__init__(
         model_dir=model_dir,
-        n_classes=0,
         weight_column_name=weight_column_name,
         linear_feature_columns=linear_feature_columns,
         linear_optimizer=linear_optimizer,
@@ -597,3 +632,44 @@ class DNNLinearCombinedRegressor(_DNNLinearCombinedBaseEstimator):
         dnn_hidden_units=dnn_hidden_units,
         dnn_activation_fn=dnn_activation_fn,
         config=config)
+
+  def predict(self, x=None, input_fn=None, batch_size=None):
+    """Returns predictions for given features.
+
+    Args:
+      x: features.
+      input_fn: Input function. If set, x must be None.
+      batch_size: Override default batch size.
+
+    Returns:
+      Numpy array of predicted classes or regression values.
+    """
+    predictions = self._infer_model(x=x,
+                                    input_fn=input_fn,
+                                    batch_size=batch_size)
+    return predictions
+
+  def _loss_vec(self, logits, target):
+    return math_ops.square(logits - math_ops.to_float(target))
+
+  def _logits_to_predictions(self, logits, proba=False):
+    return array_ops.reshape(logits, [-1])
+
+  def _num_label_columns(self):
+    return 1
+
+  def _get_eval_ops(self, features, targets, metrics=None):
+    """See base class."""
+    logits = self._logits(features)
+    result = {"loss": metrics_lib.streaming_mean(self._loss(
+        logits, targets,
+        weight_tensor=self._get_weight_tensor(features)))}
+
+    # Adding default metrics
+    if metrics:
+      predictions = self._logits_to_predictions(logits, proba=False)
+      result.update(self._run_metrics(predictions, targets, metrics,
+                                      self._get_weight_tensor(features)))
+
+    return result
+