Repository: sql-machine-learning/elasticdl

Issue:
push embedding table info and gradients in worker



Discussion:

Resolving patch:
diff --git a/elasticdl/python/common/hash_utils.py b/elasticdl/python/common/hash_utils.py
--- a/elasticdl/python/common/hash_utils.py
+++ b/elasticdl/python/common/hash_utils.py
@@ -8,3 +8,42 @@ def string_to_id(name, bucket_num):
 
 def int_to_id(number, bucket_num):
     return number % bucket_num
+
+
+def scatter_embedding_vector(values, indices, num):
+    """
+    Scatter embedding vectors to different parameter servers.
+    There are two steps to process <id, embedding vector> pairs:
+    1. hash the item id to a parameter server id
+    2. put corresponding item embedding vector to the parameter server id
+
+    For example, we scatter following embedding vectors into two parameter
+    servers:
+        values = np.array([[1, 2], [3, 4], [5, 6]])
+        indices = np.array([8, 1, 7])
+
+    ID 8 will be hashed to parameter server 0, ID 1 and ID 7 will be hashed to
+    parameter server 1. So, parameter server 0 has embedding vectors
+    np.array([[1, 2]]), parameter server 1 has embedding vectors,
+    np.array([[3, 4], [5, 6]).
+
+    The function return a dictionary:
+    {
+        0: (np.array([[1, 2]]), [8]),
+        1: (np.array([[3, 4], [5, 6]]), [1, 7])
+    }
+    """
+    ps_ids = {}
+    indices_list = indices.tolist()
+    for i, item_id in enumerate(indices_list):
+        ps_id = int_to_id(item_id, num)
+        if ps_id not in ps_ids:
+            ps_ids[ps_id] = [(i, item_id)]
+        else:
+            ps_ids[ps_id].append((i, item_id))
+    results = {}
+    for ps_id, i_item_id in ps_ids.items():
+        i = [v[0] for v in i_item_id]
+        item_id = [v[1] for v in i_item_id]
+        results[ps_id] = (values[i, :], item_id)
+    return results
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
@@ -11,7 +11,11 @@
     Mode,
     SaveModelConfig,
 )
-from elasticdl.python.common.hash_utils import int_to_id, string_to_id
+from elasticdl.python.common.hash_utils import (
+    int_to_id,
+    scatter_embedding_vector,
+    string_to_id,
+)
 from elasticdl.python.common.log_utils import default_logger as logger
 from elasticdl.python.common.model_handler import ModelHandler
 from elasticdl.python.common.model_utils import (
@@ -289,12 +293,22 @@ def report_variable_to_ps(self):
                 ps_vars[ps_id].append(v)
 
         # TODO: call `push_model` in parallel
-        for ps_id, vars in ps_vars.items():
+        for ps_id in range(len(self._ps_stubs)):
             model = elasticdl_pb2.Model()
-            for var in vars:
-                emplace_tensor_pb_from_ndarray(
-                    model.param, var.numpy(), name=var.name
-                )
+            if ps_id in ps_vars:
+                vars = ps_vars[ps_id]
+                for var in vars:
+                    emplace_tensor_pb_from_ndarray(
+                        model.param, var.numpy(), name=var.name
+                    )
+            if self._embedding_layers:
+                embedding_infos = elasticdl_pb2.EmbeddingTableInfo
+                for layer in self._embedding_layers:
+                    embedding_info = embedding_infos.add()
+                    embedding_info.name = layer.name
+                    embedding_info.dim = layer.output_dim
+                    embedding_info.initializer = layer.embeddings_initializer
+
             self._ps_stubs[ps_id].push_model(model)
 
     def report_variable(self):
@@ -371,14 +385,58 @@ def report_gradient_to_ps(self, grads):
             else:
                 ps_grads[ps_id].append((g, v.name))
 
+        if self._embedding_layers:
+            edl_embedding_grads = grads[non_embed_vars_n:]
+            bet_number = 0
+            for layer in self._embedding_layers:
+                bet_number += len(layer.embedding_and_ids)
+            if len(edl_embedding_grads) != bet_number:
+                raise ValueError(
+                    "elasticdl.layers.embedding related gradient number %d "
+                    "does not match the number of its output tensor %d."
+                    % (len(edl_embedding_grads), bet_number)
+                )
+
+            grad_accum_iter = 0
+            for layer in self._embedding_layers:
+                g_values = None
+                g_indices = None
+                for _, ids in layer.embedding_and_ids:
+                    grad = edl_embedding_grads[grad_accum_iter]
+                    grad_accum_iter += 1
+                    # ElasticDL embedding layer with Sparse Gradients
+                    if isinstance(grad, tf.IndexedSlices):
+                        grad = grad.values
+                    if g_values is not None:
+                        g_values = tf.concat([g_values, grad], axis=0)
+                        g_indices = tf.concat([g_indices, ids], axis=0)
+                    else:
+                        g_values = grad
+                        g_indices = ids
+
+                results = scatter_embedding_vector(
+                    g_values.numpy(), g_indices.numpy()
+                )
+
         # TODO: call `push_gradient` in parallel
-        for ps_id, grads in ps_grads.items():
+        for ps_id in range(len(self._ps_stubs)):
             req = elasticdl_pb2.PushGradientRequest()
-            for g, name in grads:
-                emplace_tensor_pb_from_ndarray(req.gradients, g, name=name)
+            if ps_id in ps_grads:
+                for g, name in ps_grads[ps_id]:
+                    emplace_tensor_pb_from_ndarray(req.gradients, g, name=name)
+
+            if self._embedding_layers:
+                if ps_id in results:
+                    for (gv, gi) in results[ps_id]:
+                        emplace_tensor_pb_from_ndarray(
+                            req.gradients,
+                            values=gv,
+                            indices=gi,
+                            name=layer.name,
+                        )
+
             req.model_version = self._model_version
             res = self._ps_stubs[ps_id].push_gradient(req)
-
         # TODO: choose the last response temporarily
         return res.accepted, res.model_version
 


Test patch:
diff --git a/elasticdl/python/tests/hash_utils_test.py b/elasticdl/python/tests/hash_utils_test.py
new file mode 100644
--- /dev/null
+++ b/elasticdl/python/tests/hash_utils_test.py
@@ -0,0 +1,45 @@
+import unittest
+
+import numpy as np
+
+from elasticdl.python.common.hash_utils import (
+    int_to_id,
+    scatter_embedding_vector,
+)
+
+
+class HashUtilTest(unittest.TestCase):
+    def test_scatter_embedding_vector(self):
+        vectors = np.array([[1, 2], [3, 4], [5, 6], [1, 7], [3, 9]])
+        indices = np.array([0, 1, 2, 3, 4])
+        num = 2
+
+        expected_results = {}
+        for i, item_id in enumerate(indices):
+            ps_id = int_to_id(item_id, num)
+            if ps_id not in expected_results:
+                item_list = [item_id]
+                expected_results[ps_id] = [
+                    np.expand_dims(vectors[i, :], axis=0),
+                    item_list,
+                ]
+            else:
+                expected_results[ps_id][0] = np.concatenate(
+                    (
+                        expected_results[ps_id][0],
+                        np.expand_dims(vectors[i, :], axis=0),
+                    ),
+                    axis=0,
+                )
+                expected_results[ps_id][1].append(item_id)
+        results = scatter_embedding_vector(vectors, indices, num)
+
+        for ps_id in range(num):
+            np.testing.assert_array_equal(
+                results[ps_id][0], expected_results[ps_id][0]
+            )
+            self.assertListEqual(results[ps_id][1], expected_results[ps_id][1])
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/elasticdl/python/tests/worker_and_PS_interaction_test.py b/elasticdl/python/tests/worker_ps_interaction_test.py
similarity index 99%
rename from elasticdl/python/tests/worker_and_PS_interaction_test.py
rename to elasticdl/python/tests/worker_ps_interaction_test.py
--- a/elasticdl/python/tests/worker_and_PS_interaction_test.py
+++ b/elasticdl/python/tests/worker_ps_interaction_test.py
@@ -26,7 +26,7 @@ def random_batch(batch_size):
     return images, labels
 
 
-class WorkerMNISTTest(unittest.TestCase):
+class WorkerPSInteractionTest(unittest.TestCase):
     def setUp(self):
         self._model_zoo_path = os.path.join(
             os.path.dirname(os.path.realpath(__file__)), "../../../model_zoo"