[Embedding] Fix build save graph bug when creating partitioned EmbeddingVariable in feature_column API. (#521)

Author: candyzone

tensorflow/python/feature_column/feature_column.py

@@ -155,6 +155,7 @@
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import nn_ops
 from tensorflow.python.ops import parsing_ops
+from tensorflow.python.ops import partitioned_variables
 from tensorflow.python.ops import resource_variable_ops
 from tensorflow.python.ops import sparse_ops
 from tensorflow.python.ops import string_ops

tensorflow/python/feature_column/feature_column_v2_test.py

@@ -7670,6 +7670,26 @@ def testEmbeddingVariableForSharedEmbeddingColumns(self):
         for j in range(3):
           self.assertAlmostEqual(emb_r[i][j], emb_right[i][j])
 
+  @test_util.run_deprecated_v1
+  def testEmbeddingVariableForSharedEmbeddingColumnsWithPartitionNum(self):
+    columns_list=[]
+    columns_list.append(fc.categorical_column_with_embedding("col_emb", dtype=dtypes.string, partition_num=4))
+    W = fc.shared_embedding_columns(columns_list,
+            dimension=3,
+            initializer=init_ops.ones_initializer(dtypes.float32),
+            shared_embedding_collection_name="xxxxx_shared")
+
+    ids={}
+    ids["col_emb"] = sparse_tensor.SparseTensor(indices=[[0,0],[1,0],[2,0],[3,0],[4,0]], values=["aaaa","bbbbb","ccc","4nn","5b"], dense_shape=[5, 5])
+    emb = fc_old.input_layer(ids, W)
+    fun = math_ops.multiply(emb, 2.0, name='multiply')
+    loss = math_ops.reduce_sum(fun, name='reduce_sum')
+    opt = ftrl.FtrlOptimizer(0.1, l1_regularization_strength=2.0, l2_regularization_strength=0.00001)
+    g_v = opt.compute_gradients(loss)
+    train_op = opt.apply_gradients(g_v)
+    init = variables_lib.global_variables_initializer()
+    saver = saver_module.Saver()
+
   @test_util.run_deprecated_v1
   def test_transform_feature(self):
     a = fc.categorical_column_with_identity(key='aaa', num_buckets=3)

tensorflow/python/ops/embedding_variable_ops_test.py

@@ -2281,5 +2281,31 @@ def runTestAdagrad(self, var, g):
       for j in range(0, 3):
         self.assertEqual(emb1[i][j], emb2[i][j])
 
+  def testEmbeddingVariableForContirbFeatureColumnWithPartitionNum(self):
+    print("testEmbeddingVariableForContirbFeatureColumnWithPartitionNum")
+    checkpoint_directory = self.get_temp_dir()
+    evict = variables.L2WeightEvict(l2_weight_threshold=0.9)
+    columns = feature_column.sparse_column_with_embedding(
+                                        column_name="col_emb",
+                                        dtype=dtypes.int64,
+                                        partition_num = 4)
+    W = feature_column.embedding_column(sparse_id_column=columns,
+            dimension=3,
+            initializer=init_ops.ones_initializer(dtypes.float32),
+            combiner="mean")
+    ids = {}
+    ids["col_emb"] = sparse_tensor.SparseTensor(
+                      indices=[[0,0],[1,0],[2,0],[3,0],[4,0],[5,0]],
+                      values=math_ops.cast([0,0,0,1,1,2], dtypes.int64),
+                      dense_shape=[6, 1])
+    emb= feature_column_ops.input_from_feature_columns(
+           columns_to_tensors=ids, feature_columns=[W])
+    fun = math_ops.multiply(emb, 2.0, name='multiply')
+    loss = math_ops.reduce_sum(fun, name='reduce_sum')
+    opt = ftrl.FtrlOptimizer(0.1, l1_regularization_strength=2.0, l2_regularization_strength=0.00001)
+    g_v = opt.compute_gradients(loss)
+    train_op = opt.apply_gradients(g_v)
+    saver = saver_module.Saver()
+
 if __name__ == "__main__":
   googletest.main()

tensorflow/python/training/saver.py

@@ -701,7 +701,8 @@ def _build_internal(self,
               # not need it, but we'll try looking it up on MetaGraph restore
               # since it's in a collection.
               params=list(element)
-              if not isinstance(params[0], kv_variable_ops.DynamicEmbeddingVariable):
+              if not isinstance(params[0], kv_variable_ops.DynamicEmbeddingVariable) and \
+                 not isinstance(params[0], kv_variable_ops.EmbeddingVariable):
                 element.as_tensor()
       return saver_pb2.SaverDef(
           filename_tensor_name=filename_tensor.name,