Support weighted sparse column in sdca optimizer.

tensorflower-gardener · tensorflower-gardener · commit a9e876827435 · 2016-07-11T09:32:43.000-07:00
Change: 127093435
diff --git a/tensorflow/contrib/learn/python/learn/estimators/linear_test.py b/tensorflow/contrib/learn/python/learn/estimators/linear_test.py
@@ -203,6 +203,33 @@ def input_fn():
     scores = classifier.evaluate(input_fn=input_fn, steps=2)
     self.assertGreater(scores['accuracy'], 0.9)
 
+  def testSdcaOptimizerWeightedSparseFeatures(self):
+    """LinearClasssifier with SDCAOptimizer and weighted sparse features."""
+
+    def input_fn():
+      return {
+          'example_id': tf.constant(['1', '2', '3']),
+          'price': tf.SparseTensor(values=[2., 3., 1.],
+                                   indices=[[0, 0], [1, 0], [2, 0]],
+                                   shape=[3, 5]),
+          'country': tf.SparseTensor(values=['IT', 'US', 'GB'],
+                                     indices=[[0, 0], [1, 0], [2, 0]],
+                                     shape=[3, 5])
+      }, tf.constant([[1], [0], [1]])
+
+    country = tf.contrib.layers.sparse_column_with_hash_bucket(
+        'country', hash_bucket_size=5)
+    country_weighted_by_price = tf.contrib.layers.weighted_sparse_column(
+        country, 'price')
+    sdca_optimizer = tf.contrib.learn.SDCAOptimizer(
+        example_id_column='example_id')
+    classifier = tf.contrib.learn.LinearClassifier(
+        feature_columns=[country_weighted_by_price],
+        optimizer=sdca_optimizer)
+    classifier.fit(input_fn=input_fn, steps=50)
+    scores = classifier.evaluate(input_fn=input_fn, steps=2)
+    self.assertGreater(scores['accuracy'], 0.9)
+
   def testSdcaOptimizerCrossedFeatures(self):
     """Tests LinearClasssifier with SDCAOptimizer and crossed features."""
 
diff --git a/tensorflow/contrib/learn/python/learn/estimators/sdca_optimizer.py b/tensorflow/contrib/learn/python/learn/estimators/sdca_optimizer.py
@@ -124,6 +124,17 @@ def _training_examples_and_variables():
                                                            column.length)
           sparse_features.append(math_ops.to_float(sparse_features_tensor))
           sparse_features_weights.append(columns_to_variables[column][0])
+        elif isinstance(
+            column,
+            layers.feature_column._WeightedSparseColumn):  # pylint: disable=protected-access
+          id_tensor = column.id_tensor(transformed_tensor)
+          weight_tensor = column.weight_tensor(transformed_tensor)
+          sparse_features_tensor = sparse_ops.sparse_merge(
+              id_tensor, weight_tensor, column.length,
+              name="{}_sparse_merge".format(column.name))
+          sparse_features.append(math_ops.to_float(
+              sparse_features_tensor, name="{}_to_float".format(column.name)))
+          sparse_features_weights.append(columns_to_variables[column][0])
         else:
           raise ValueError("SDCAOptimizer does not support column type %s." %
                            type(column).__name__)
