[Embedding] Add group_strategy to control parallelism of group_embedding.

Signed-off-by: JunqiHu <[email protected]>

Author: Mesilenceki

Diff for: tensorflow/python/BUILD

@@ -968,7 +968,7 @@ py_library(
         ":type_spec",
         ":util",
         ":versions",
-        ":group_embedding_ops_utils",
+        ":group_embedding_types",
         "//tensorflow/core:protos_all_py",
         "//tensorflow/python/eager:context",
         "//tensorflow/python/eager:core",
@@ -985,8 +985,8 @@ py_library(
 )
 
 py_library(
-    name = "group_embedding_ops_utils",
-    srcs = ["framework/group_embedding_ops_utils.py"],
+    name = "group_embedding_types",
+    srcs = ["framework/group_embedding_types.py"],
     srcs_version = "PY2AND3",
     deps = [
         ":pywrap_tensorflow",

Diff for: tensorflow/python/feature_column/feature_column_v2.py

@@ -131,6 +131,7 @@
 import contextlib
 import math
 
+import sys
 import numpy as np
 import six
 import json
@@ -4235,21 +4236,21 @@ def _from_config(cls, config, custom_objects=None, columns_by_name=None):
 
 @tf_export('feature_column.group_embedding_column_scope')
 @contextlib.contextmanager
-def group_embedding_column_scope(name=''):
+def group_embedding_column_scope(name='', params_num_per_group=sys.maxsize):
   global_group_embedding_scope = group_embedding_column._global_group_embedding_scope_list()
   group_id = group_embedding_column._current_group_id()
   if name == '':
     name = "group_embedding_column_scope_{}".format(group_id)
     group_id +=1
   else:
      name = "group_embedding_column_scope_{}".format(name)
-  fusion_embedding_scope = GroupEmbeddingScope(name)
+  fusion_embedding_scope = GroupEmbeddingScope(name, params_num_per_group)
   global_group_embedding_scope.append(fusion_embedding_scope)
   yield global_group_embedding_scope 
 
 class GroupEmbeddingScope(group_embedding_column.GroupEmbeddingScopeBase):
-  def __init__(self, name=None):
-    super(GroupEmbeddingScope, self).__init__(name=name)
+  def __init__(self, name=None, params_num_per_group=sys.maxsize):
+    super(GroupEmbeddingScope, self).__init__(name=name, params_num_per_group=params_num_per_group)
 
   def add_column(self, embedding_column):
     VALID_EMBEDDING_COLUMN_TYPES = (
@@ -4288,7 +4289,8 @@ def _get_dense_tensor(self, filter_ec, inputs,
         embedding_weights.append(embedding_weight)
 
     output_tensors.extend(embedding_ops.group_embedding_lookup_sparse(
-                              embedding_weights, sp_ids, combiners, is_sequence=is_sequence))
+                              embedding_weights, sp_ids, combiners, 
+                              is_sequence=is_sequence, params_num_per_group=self.params_num_per_group))
     return output_tensors, sequence_lengths
 
 class EmbeddingColumn(

Diff for: tensorflow/python/feature_column/group_embedding_column.py

@@ -2,6 +2,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import sys
+
 _global_fusion_embedding_scope = []
 _group_id = 0
 _group_embedding_tensor = dict()
@@ -40,8 +42,9 @@ def _current_group_id():
     return _group_id
 
 class GroupEmbeddingScopeBase(object):
-    def __init__(self, name=None):
+    def __init__(self, name=None, params_num_per_group=sys.maxsize):
         self.name = name
+        self.params_num_per_group = params_num_per_group
         self.embedding_columns = []
 
     def add_column(self, embedding_column):

Diff for: tensorflow/python/framework/config.py

@@ -21,7 +21,7 @@
 from tensorflow.python.eager import context
 from tensorflow.python import _pywrap_tensor_float_32_execution
 from tensorflow.python.util.tf_export import tf_export
-from tensorflow.python.framework import group_embedding_ops_utils
+from tensorflow.python.framework import group_embedding_types
 
 @tf_export('config.experimental.tensor_float_32_execution_enabled')
 def tensor_float_32_execution_enabled():
@@ -646,7 +646,7 @@ def enable_distributed_strategy(strategy="collective"):
     try:
       from sparse_operation_kit import experiment as sok
       sok.init()
-      group_embedding_ops_utils.set_group_lookup_strategy(strategy)
+      group_embedding_types.set_group_lookup_strategy(strategy)
     except:
       raise ImportError("While param `strategy` in enable_distributed_strategy"
                         "is given `collective`, sok module initialize error,"

Diff for: tensorflow/python/framework/group_embedding_ops_utils.py renamed to tensorflow/python/framework/group_embedding_types.py

@@ -21,26 +21,26 @@
 from enum import Enum, unique
 
 @unique
-class STRATEGY(Enum):
+class DistStrategy(Enum):
   COLLECTIVE = "collective"
   DISTRIBUTED = "ps"
   LOCALIZED = "localized"
   UNKNOWN = "unknown"
 
-_group_lookup_strategy = STRATEGY.LOCALIZED
+_group_lookup_strategy = DistStrategy.LOCALIZED
 
 def set_group_lookup_strategy(strategy):
   def str_to_strategy(strategy):
     if strategy == "collective":
-      return STRATEGY.COLLECTIVE
+      return DistStrategy.COLLECTIVE
     elif strategy == "ps":
-      return STRATEGY.DISTRIBUTED
+      return DistStrategy.DISTRIBUTED
     elif strategy == "localized":
-      return STRATEGY.LOCALIZED
+      return DistStrategy.LOCALIZED
 
   global _group_lookup_strategy
   _group_lookup_strategy = str_to_strategy(strategy)
 
 def get_group_lookup_strategy():
   global _group_lookup_strategy
-  return _group_lookup_strategy
+  return _group_lookup_strategy