Repository: sql-machine-learning/elasticdl

Issue:
elasticdl.layers.Embedding does not need worker.
It only needs a lookup function to lookup embeddings and a tape for automatic gradient computation.


Discussion:

Resolving patch:
diff --git a/elasticdl/python/elasticdl/layers/embedding.py b/elasticdl/python/elasticdl/layers/embedding.py
--- a/elasticdl/python/elasticdl/layers/embedding.py
+++ b/elasticdl/python/elasticdl/layers/embedding.py
@@ -45,7 +45,7 @@ def __init__(
         self.input_length = input_length
         self.combiner = combiner
         self.tape = None
-        self.worker = None
+        self.lookup_func = None
         self.bet_ids_pair = []
 
     @tf_utils.shape_type_conversion
@@ -87,7 +87,7 @@ def get_key(name_list):
         return "-".join(map(str, name_list))
 
     def lookup_embedding(self, unique_ids):
-        batch_embedding = self.worker.lookup_embedding(
+        batch_embedding = self.lookup_func(
             unique_ids.numpy(),
             self._name,
             self.embedding_initializer,
@@ -161,10 +161,19 @@ def compute_mask(self, inputs, mask=None):
 
     def reset(self):
         self.bet_ids_pair = []
-        self.tape = None
 
     def set_tape(self, tape):
         self.tape = tape
 
-    def set_worker(self, worker):
-        self.worker = worker
+    def set_lookup_func(self, lookup_func):
+        """
+        lookup_func args
+          ids: id list to lookup
+          layer_name: layer name
+          initializer: intializer method name for unkown embeddings
+          output_dim: the embedding vector dimension
+        return:
+          embeddings corresponding to the ids, in numpy ndarray type
+          with a shape of (len(ids), output_dim)
+        """
+        self.lookup_func = lookup_func
diff --git a/elasticdl/python/worker/worker.py b/elasticdl/python/worker/worker.py
--- a/elasticdl/python/worker/worker.py
+++ b/elasticdl/python/worker/worker.py
@@ -61,6 +61,12 @@ def __init__(
         self._model = load_model_from_module(
             model_def, model_module, model_params
         )
+        self._tape = None
+        if (
+            self._job_type == JobType.TRAINING_WITH_EVALUATION
+            or self._job_type == JobType.TRAINING_ONLY
+        ):
+            self._tape = tf.GradientTape(persistent=True)
         self._init_embedding_layer()
         self._var_created = self._model.built
         self._dataset_fn = model_module[dataset_fn]
@@ -100,7 +106,8 @@ def _init_embedding_layer(self):
         """
         self._embedding_layers = find_layer(self._model, Embedding)
         for layer in self._embedding_layers:
-            layer.set_worker(self)
+            layer.set_lookup_func(self.lookup_embedding)
+        self._set_tape_for_embedding(self._tape)
         if self._embedding_layers:
             # TODO check that Redis IP/PORT is set
             pass
@@ -321,14 +328,13 @@ def training_process_with_acceleration(self, features, labels):
         return self.training_process_eagerly(features, labels)
 
     def training_process_eagerly(self, features, labels):
-        with tf.GradientTape() as tape:
-            self._set_tape_for_embedding(tape)
+        with self._tape:
             outputs = self._model.call(features, training=True)
             loss = self._loss(outputs, labels)
             # Add regularization loss if any
             if self._model.losses:
                 loss += tf.math.add_n(self._model.losses)
-        grads = tape.gradient(loss, self.get_trainable_items())
+        grads = self._tape.gradient(loss, self.get_trainable_items())
         return loss, grads
 
     @tf.function
@@ -349,8 +355,15 @@ def _run_training_task(self, features, labels):
         return accepted, min_model_version, loss
 
     def _run_evaluation_task(self, features, labels):
+        # reset tape to None for evaluation
+        if self._tape:
+            self._set_tape_for_embedding(None)
         evaluation_metrics = self.evaluation_process(features, labels)
-        return self.report_evaluation_metrics(evaluation_metrics)
+        accepted, _ = self.report_evaluation_metrics(evaluation_metrics)
+        # set tape back for training
+        if self._tape:
+            self._set_tape_for_embedding(self._tape)
+        return accepted
 
     def _run_prediction_task(self, features):
         predictions = self.predict_process(features)
@@ -368,7 +381,7 @@ def _process_minibatch(
                         self.get_model(0, elasticdl_pb2.MINIMUM)
                 elif self._model_version != min_model_version:
                     self.get_model(min_model_version, elasticdl_pb2.FIXED)
-                accepted, _ = self._run_evaluation_task(features, labels)
+                accepted = self._run_evaluation_task(features, labels)
                 if accepted:
                     break
             elif task_type == elasticdl_pb2.TRAINING:


Test patch:
diff --git a/elasticdl/python/tests/layer_test.py b/elasticdl/python/tests/layer_test.py
--- a/elasticdl/python/tests/layer_test.py
+++ b/elasticdl/python/tests/layer_test.py
@@ -84,7 +84,7 @@ def create_embedding_layer(
         mask_zero=mask_zero,
     )
     worker = mock_worker(embedding_size, output_dim)
-    layer.set_worker(worker)
+    layer.set_lookup_func(worker.lookup_embedding)
     return layer