Repository: sql-machine-learning/elasticdl

Issue:
embedding table Python data structure



Discussion:

Resolving patch:
diff --git a/elasticdl/python/ps/__init__.py b/elasticdl/python/ps/__init__.py
new file mode 100644
diff --git a/elasticdl/python/ps/embedding_table.py b/elasticdl/python/ps/embedding_table.py
new file mode 100644
--- /dev/null
+++ b/elasticdl/python/ps/embedding_table.py
@@ -0,0 +1,51 @@
+import numpy as np
+import tensorflow as tf
+
+
+class EmbeddingTable(object):
+    """
+    EmbeddingTable is used to store embedding parameters of an embedding
+    layer. The name of an embedding table is actually the embedding layer
+    name. It uses a dictionary to store embedding vectors, the key is the
+    item id, the value is a 1-D numpy.ndarray.
+
+    Embedding vectors are lazily initialized in parameter server.
+    EmbeddingTable also has dim and initializer fields. Inside the get
+    interface of EmbeddingTable, if the id is not in the embedding_vectors
+    dictionary, the corresponding value will be initialized.
+    """
+
+    def __init__(self, name, dim=None, initializer=None):
+        self.name = name
+        self.dim = dim
+        self.initializer = initializer
+        self.embedding_vectors = {}
+
+    def get(self, indices):
+        if len(indices) == 0:
+            return None
+        values = []
+        for i in indices:
+            value = self.embedding_vectors.get(i, None)
+            if value is None:
+                init = tf.keras.initializers.get(self.initializer)
+                value = init(shape=(self.dim,)).numpy()
+                self.embedding_vectors[i] = value
+            values.append(value)
+        return np.stack(values)
+
+    def set(self, indices, values):
+        # TODO(qijun) need to add a RWLock in Sync-SGD
+        for index, i in enumerate(indices):
+            embedding_vector = values[index, :]
+            self.embedding_vectors[i] = embedding_vector
+
+    def clear(self):
+        self.embedding_vectors.clear()
+
+
+def create_embedding_table(embedding_table_info_pb):
+    name = embedding_table_info_pb.name
+    dim = embedding_table_info_pb.dim
+    initializer = embedding_table_info_pb.initializer
+    return EmbeddingTable(name, dim, initializer)


Test patch:
diff --git a/elasticdl/python/tests/embedding_table_test.py b/elasticdl/python/tests/embedding_table_test.py
new file mode 100644
--- /dev/null
+++ b/elasticdl/python/tests/embedding_table_test.py
@@ -0,0 +1,61 @@
+import unittest
+
+import numpy as np
+
+from elasticdl.proto.elasticdl_pb2 import EmbeddingTableInfo
+from elasticdl.python.ps.embedding_table import (
+    EmbeddingTable,
+    create_embedding_table,
+)
+
+
+class EmbeddingTableTest(unittest.TestCase):
+    def setUp(self):
+        self.name = "embedding_1"
+        self.dim = 10
+        self.initializer = "uniform"
+        self.table = EmbeddingTable(self.name, self.dim, self.initializer)
+
+    def test_embedding_table_init(self):
+        self.assertIsNotNone(self.table)
+        self.assertEqual(self.table.name, self.name)
+        self.assertEqual(self.table.dim, self.dim)
+        self.assertEqual(self.table.initializer, self.initializer)
+
+    def test_embedding_table_get(self):
+        self.table.clear()
+        indices = [0, 3, 7]
+        res = self.table.get(indices)
+        self.assertTupleEqual(res.shape, (3, 10))
+
+        res = self.table.get([])
+        self.assertIsNone(res)
+
+        self.table.get([0, 3, 8])
+        self.assertEqual(len(self.table.embedding_vectors), 4)
+
+    def test_embedding_table_set(self):
+        self.table.clear()
+        indices = [0, 1, 4]
+        x = len(indices)
+        values = np.random.uniform(size=x * self.dim).reshape((x, self.dim))
+        self.table.set(indices, values)
+
+        row0 = self.table.get([0])
+        row1 = self.table.get([1])
+        row4 = self.table.get([4])
+
+        rows = [row0, row1, row4]
+        rows = np.concatenate(rows)
+        np.testing.assert_array_equal(rows, values)
+
+    def test_create_embedding_table(self):
+        embedding_pb = EmbeddingTableInfo()
+        embedding_pb.name = self.name
+        embedding_pb.dim = self.dim
+        embedding_pb.initializer = self.initializer
+        table = create_embedding_table(embedding_pb)
+        self.assertIsNotNone(table)
+        self.assertEqual(table.name, self.name)
+        self.assertEqual(table.initializer, self.initializer)
+        self.assertEqual(table.dim, self.dim)