Repository: sql-machine-learning/elasticdl

Issue:
Worker assigns self to EdlEmbedding layers and sanity check
Read data from Redis.


Discussion:

Resolving patch:
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
@@ -6,6 +6,7 @@
 from elasticdl.python.common.constants import JobType, Mode
 from elasticdl.python.common.log_util import default_logger as logger
 from elasticdl.python.common.model_helper import (
+    find_layer,
     load_model_from_module,
     load_module,
 )
@@ -13,6 +14,7 @@
     ndarray_to_tensor,
     tensor_to_ndarray,
 )
+from elasticdl.python.elasticdl.layers.embedding import Embedding
 from elasticdl.python.worker.prediction_outputs_processor import (
     BasePredictionOutputsProcessor,
 )
@@ -56,8 +58,7 @@ def __init__(
         self._model = load_model_from_module(
             model_def, model_module, model_params
         )
-        # TODO: call find_layer and get Embedding layers
-        self._embedding_layers = []
+        self._init_embedding_layer()
         self._var_created = self._model.built
         self._dataset_fn = model_module[dataset_fn]
         self._opt_fn = model_module[optimizer]
@@ -89,6 +90,25 @@ def __init__(
         else:
             self._prediction_outputs_processor = None
 
+    def _init_embedding_layer(self):
+        """
+        Init elasticdl.layers.embedding layer list and assign worker to them
+        """
+        self._embedding_layers = find_layer(self._model, Embedding)
+        for layer in self._embedding_layers:
+            layer.set_worker(self)
+        if self._embedding_layers:
+            # TODO check that Redis IP/PORT is set
+            pass
+
+    def _set_tape_for_embedding(self, tape):
+        for layer in self._embedding_layers:
+            layer.set_tape(tape)
+
+    def _reset_embedding(self):
+        for layer in self._embedding_layers:
+            layer.reset()
+
     def get_task(self):
         """
         get task from master
@@ -238,6 +258,7 @@ def get_trainable_items(self):
     @tf.function
     def training_process(self, features, labels):
         with tf.GradientTape() as tape:
+            self._set_tape_for_embedding(tape)
             outputs = self._model.call(features, training=True)
             loss = self._loss(outputs, labels)
             # Add regularization loss if any
@@ -260,6 +281,7 @@ def predict_process(self, features):
     def _run_training_task(self, features, labels):
         loss, grads = self.training_process(features, labels)
         accepted, min_model_version = self.report_gradient(grads)
+        self._reset_embedding()
         return accepted, min_model_version, loss
 
     def _run_evaluation_task(self, features, labels):


Test patch:
diff --git a/elasticdl/python/tests/embedding_test_module.py b/elasticdl/python/tests/embedding_test_module.py
new file mode 100644
--- /dev/null
+++ b/elasticdl/python/tests/embedding_test_module.py
@@ -0,0 +1,50 @@
+import tensorflow as tf
+from tensorflow.keras.layers import Concatenate, Dense, Flatten
+
+from elasticdl.python.elasticdl.layers.embedding import Embedding
+
+
+class CustomModel(tf.keras.Model):
+    def __init__(self, embedding_dim=16):
+        super(CustomModel, self).__init__(name="embedding_test_model")
+        self.embedding_dim = embedding_dim
+        self.embedding_1 = Embedding(embedding_dim)
+        self.embedding_2 = Embedding(embedding_dim)
+        self.concat = Concatenate()
+        self.dense = Dense(1, input_shape=(embedding_dim * 3,))
+        self.flatten = Flatten()
+
+    def call(self, inputs, training=False):
+        f1 = self.embedding_1(inputs["f1"])
+        f2 = self.embedding_1(inputs["f2"])
+        f3 = self.embedding_2(inputs["f3"])
+        x = self.concat([f1, f2, f3])
+        x = self.dense(x)
+        return self.flatten(x)
+
+
+def loss(predictions, labels):
+    return tf.reduce_mean(tf.square(predictions - labels))
+
+
+def dataset_fn(dataset, training=True):
+    def _parse_data(record):
+        feature_description = {
+            "f1": tf.io.FixedLenFeature([1], tf.int64),
+            "f2": tf.io.FixedLenFeature([1], tf.int64),
+            "f3": tf.io.FixedLenFeature([1], tf.int64),
+            "label": tf.io.FixedLenFeature([1], tf.int64),
+        }
+        r = tf.io.parse_single_example(record, feature_description)
+        return {"f1": r["f1"], "f2": r["f2"], "f3": r["f3"]}, r["label"]
+
+    dataset = dataset.map(_parse_data)
+    return dataset
+
+
+def optimizer(lr=0.1):
+    return tf.optimizers.SGD(lr)
+
+
+def eval_metrics_fn(predictions, labels):
+    return {"mse": tf.reduce_mean(tf.square(predictions - labels))}
diff --git a/elasticdl/python/tests/worker_test.py b/elasticdl/python/tests/worker_test.py
--- a/elasticdl/python/tests/worker_test.py
+++ b/elasticdl/python/tests/worker_test.py
@@ -19,6 +19,10 @@
 _model_file = get_model_file(
     os.path.dirname(os.path.realpath(__file__)), "test_module.custom_model"
 )
+_embedding_model_file = get_model_file(
+    os.path.dirname(os.path.realpath(__file__)),
+    "embedding_test_module.CustomModel",
+)
 m = load_module(_model_file).__dict__
 
 
@@ -124,6 +128,17 @@ def test_distributed_train_tf_example(self):
     def test_distributed_evaluate_tf_example(self):
         self.distributed_train_and_evaluate(training=False)
 
+    def test_embedding_layer(self):
+        worker = Worker(
+            1,
+            JobType.TRAINING_ONLY,
+            32,
+            _embedding_model_file,
+            model_def="embedding_test_module.CustomModel",
+            channel=None,
+        )
+        self.assertTrue(len(worker._embedding_layers) == 2)
+
 
 if __name__ == "__main__":
     unittest.main()