[Graph] Support asynchronous embedding lookup. (#376)

Author: JackMoriarty

tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column.py

@@ -131,7 +131,9 @@ def sequence_input_layer(
     fc._verify_static_batch_size_equality(sequence_lengths, ordered_columns)
     sequence_length = _assert_all_equal_and_return(sequence_lengths)
 
-    return array_ops.concat(output_tensors, -1), sequence_length
+    concat_result = array_ops.concat(output_tensors, -1)
+    ops.add_to_collection(ops.GraphKeys.ASYNC_EMBEDDING_OUTPUT_TENSORS, concat_result)
+    return concat_result, sequence_length
 
 
 def concatenate_context_input(context_input, sequence_input):

tensorflow/contrib/layers/python/layers/embedding_ops.py

@@ -186,6 +186,7 @@ def safe_embedding_lookup_sparse(embedding_weights,
     final_result.set_shape(
         tensor_shape.unknown_shape(
             (original_rank_dim - 1).value).concatenate(result.get_shape()[1:]))
+    ops.add_to_collections(ops.GraphKeys.ASYNC_EMBEDDING_OUTPUT_TENSORS, final_result)
     return final_result
 
 def fused_safe_embedding_lookup_sparse(embedding_weights,

tensorflow/contrib/layers/python/layers/feature_column_ops.py

@@ -164,9 +164,9 @@ def _input_from_feature_columns(columns_to_tensors,
                              '{}, {}'.format(column.name, e, ee))
         if cols_to_outs is not None:
           cols_to_outs[column] = output_tensors[-1]
+        ops.add_to_collections(ops.GraphKeys.ASYNC_EMBEDDING_OUTPUT_TENSORS, output_tensors[-1])
     return array_ops.concat(output_tensors, output_rank - 1)
 
-
 def input_from_feature_columns(columns_to_tensors,
                                feature_columns,
                                weight_collections=None,

tensorflow/core/common_runtime/graph_execution_state.cc

@@ -763,14 +763,19 @@ Status GraphExecutionState::InitBaseGraph(std::unique_ptr<Graph>&& new_graph) {
   if (session_optimizer_options.do_smart_stage() ||
       session_optimizer_options.do_smart_stage_gpu()) {
     VLOG(2) << "RUN Graph Optimization: SmartStage";
-    std::string tn;
-    ReadStringFromEnvVar("TARGET_NODES_NAME", "", &tn);
-    std::vector<std::string> target_nodes;
-    for (std::string s : str_util::Split(tn, ';')) {
-      target_nodes.push_back(s.substr(0, s.find_last_of(':')));
+
+    if (session_optimizer_options.do_async_embedding()) {
+      VLOG(0) << "Async Embedding is enable, disable SmartStage";
+    } else {
+      std::string tn;
+      ReadStringFromEnvVar("TARGET_NODES_NAME", "", &tn);
+      std::vector<std::string> target_nodes;
+      for (std::string s : str_util::Split(tn, ';')) {
+	target_nodes.push_back(s.substr(0, s.find_last_of(':')));
+      }
+      SmartStageGraph(&new_graph, target_nodes, 
+		      session_optimizer_options.do_smart_stage_gpu());
     }
-    SmartStageGraph(&new_graph, target_nodes, 
-                    session_optimizer_options.do_smart_stage_gpu());
   }
 
   SaveStatefulNodes(new_graph.get());

tensorflow/core/protobuf/config.proto

@@ -256,6 +256,9 @@ message OptimizerOptions {
   int32 micro_batch_num = 9;
   bool do_smart_stage = 10;
   bool do_smart_stage_gpu = 11;
+  bool do_async_embedding = 12;
+  int32 async_embedding_threads_num = 13;
+  int32 async_embedding_capacity = 14;
 }
 
 message GraphOptions {

tensorflow/python/BUILD

@@ -4641,6 +4641,8 @@ py_library(
         ":util",
         ":variable_scope",
         ":variables",
+	":prefetch",
+	":prefetch_runner",
         "//tensorflow/core:protos_all_py",
         "//tensorflow/python/data/ops:dataset_ops",
         "//tensorflow/python/distribute:distribute_coordinator_context",
@@ -6257,6 +6259,20 @@ tf_py_test(
     ],
 )
 
+tf_py_test(
+    name = "async_embedding_stage_test",
+    size = "small",
+    srcs = ["training/async_embedding_stage_test.py"],
+    additional_deps = [
+        ":training",
+	":prefetch",
+	":prefetch_runner",
+	":variables",
+	":math_ops",
+	"framework",
+    ],
+)
+
 py_library(
     name = "training_util",
     srcs = ["training/training_util.py"],

tensorflow/python/feature_column/feature_column.py

@@ -218,6 +218,7 @@ def _get_logits():  # pylint: disable=missing-docstring
               scope=variable_scope.get_variable_scope().name)
         if cols_to_output_tensors is not None:
           cols_to_output_tensors[column] = output_tensor
+        ops.add_to_collections(ops.GraphKeys.ASYNC_EMBEDDING_OUTPUT_TENSORS, output_tensor)
     _verify_static_batch_size_equality(output_tensors, ordered_columns)
     return array_ops.concat(output_tensors, -1)
 

tensorflow/python/framework/ops.py

@@ -6102,6 +6102,9 @@ class GraphKeys(object):
   EV_INIT_VAR_OPS = "ev_init_var_ops"
   EV_INIT_SLOT_OPS = "ev_init_slot_ops"
 
+  # Key to collect embedding lookup output result.
+  ASYNC_EMBEDDING_OUTPUT_TENSORS = "async_embedding_output_tensors"
+
   # Key to indicate various ops.
   INIT_OP = "init_op"
   LOCAL_INIT_OP = "local_init_op"

tensorflow/python/ops/embedding_ops.py

@@ -339,6 +339,7 @@ def _embedding_lookup_and_transform(params,
       if not transform_fn:
         # If transform_fn was provided, the clip_by_norm was done above.
         ret = _clip(ret, ids, max_norm)
+      ops.add_to_collections(ops.GraphKeys.ASYNC_EMBEDDING_OUTPUT_TENSORS, ret)
       return ret
 
 
@@ -672,6 +673,7 @@ def embedding_lookup_sparse(params,
       else:
         assert False, "Unrecognized combiner"
 
+    ops.add_to_collections(ops.GraphKeys.ASYNC_EMBEDDING_OUTPUT_TENSORS, embeddings)
     return embeddings
 
 @tf_export(v1=["nn.adaptive_embedding_lookup_sparse"])
@@ -1341,6 +1343,7 @@ def safe_embedding_lookup_sparse(embedding_weights,
         tensor_shape.unknown_shape(
             (tensor_shape.Dimension(original_rank_dim) - 1).value).concatenate(
                 result.get_shape()[1:]))
+    ops.add_to_collections(ops.GraphKeys.ASYNC_EMBEDDING_OUTPUT_TENSORS, final_result)
     return final_result
 
 def fused_safe_embedding_lookup_sparse(embedding_weights,
@@ -1420,6 +1423,7 @@ def fused_safe_embedding_lookup_sparse(embedding_weights,
         tensor_shape.unknown_shape(
             (tensor_shape.Dimension(original_rank_dim) - 1).value).concatenate(
                 result.get_shape()[1:]))
+    ops.add_to_collections(ops.GraphKeys.ASYNC_EMBEDDING_OUTPUT_TENSORS, final_result)
     return final_result
 
 @tf_export("nn.safe_embedding_lookup_multi_dim")

tensorflow/python/ops/fused_embedding_ops.py

@@ -82,6 +82,7 @@ def fused_embedding_lookup_sparse(params,
       partitioned_values=partitioned_values,
       combiner=combiner, max_norm=max_norm, default_id=default_id
     )
+  ops.add_to_collections(ops.GraphKeys.ASYNC_EMBEDDING_OUTPUT_TENSORS, emb_vectors)
   return emb_vectors