Repository: sql-machine-learning/elasticdl

Issue:
EdlEmbedding.call
The forward pass according to [design doc](https://github.com/wangkuiyi/elasticdl/blob/develop/elasticdl/doc/distributed_embedding_layer_design.md#pseudocode-for-embeddingcall)


Discussion:

Resolving patch:
diff --git a/elasticdl/python/layers/embedding.py b/elasticdl/python/layers/embedding.py
new file mode 100644
--- /dev/null
+++ b/elasticdl/python/layers/embedding.py
@@ -0,0 +1,57 @@
+import tensorflow as tf
+from tensorflow.python.keras.utils import tf_utils
+
+
+class EdlEmbedding(tf.keras.layers.Layer):
+    """
+    Input: indexes for the embedding entries
+           shape is (batch_size, input_length)
+    Output: Corresponding embedding vectors of the input indexes
+            shape is (batch_size, input_length, embedding_dim)
+    Arguments:
+      embedding_dim: the dimension of the embedding vector
+      embedding_initializer: Initializer for embedding table
+    """
+
+    def __init__(self, embedding_dim, embedding_initializer="uniform"):
+        super(EdlEmbedding, self).__init__()
+        self.embedding_dim = embedding_dim
+        self.embedding_initializer = embedding_initializer
+        self.tape = None
+        self.worker = None
+        self.bet_ids_pair = []
+
+    @tf_utils.shape_type_conversion
+    def compute_output_shape(self, input_shape):
+        return input_shape + (self.embedding_dim,)
+
+    @property
+    def name(self):
+        return self._name
+
+    def call(self, input):
+        ids = tf.convert_to_tensor(input, name="embedding_ids")
+        flat_ids = tf.reshape(ids, [-1])
+        unique_ids, idx = tf.unique(flat_ids)
+        batch_embedding = self.worker.embedding_lookup(
+            unique_ids.numpy(), self._name, self.embedding_initializer
+        )
+        batch_embedding_tensor = tf.convert_to_tensor(batch_embedding)
+        if self.tape:
+            self.tape.watch(batch_embedding_tensor)
+            self.bet_ids_pair.append((batch_embedding_tensor, unique_ids))
+        outputs = tf.gather(batch_embedding_tensor, idx)
+        outputs = tf.reshape(
+            outputs, ids.get_shape().concatenate(self.embedding_dim)
+        )
+        return outputs
+
+    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


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
@@ -1,6 +1,7 @@
 import os
 import unittest
 
+import numpy as np
 import tensorflow as tf
 
 from elasticdl.python.common.model_helper import (
@@ -9,6 +10,7 @@
     load_model_from_module,
     load_module,
 )
+from elasticdl.python.layers.embedding import EdlEmbedding
 
 
 def _get_model_zoo_path():
@@ -53,5 +55,80 @@ def test_find_layer_nested(self):
             self.assertEqual(layer_num[layer_class], len(layers))
 
 
+class mock_worker:
+    def __init__(self, embedding_size, embedding_dim):
+        self.embedding_size = embedding_size
+        self.embedding_dim = embedding_dim
+        self.embedding = np.ndarray(
+            shape=(embedding_size, embedding_dim), dtype=np.float32
+        )
+        for i in range(embedding_size):
+            self.embedding[i].fill(i)
+
+    def embedding_lookup(self, ids, name, embedding_initializer):
+        values = np.take(self.embedding, ids, axis=0)
+        return values
+
+
+def create_embedding_layer(embedding_size, embedding_dim):
+    layer = EdlEmbedding(embedding_dim)
+    worker = mock_worker(embedding_size, embedding_dim)
+    layer.set_worker(worker)
+    return layer
+
+
+class EmbeddingLayerTest(unittest.TestCase):
+    def test_embedding_layer(self):
+        embedding_dim = 8
+        embedding_size = 16
+        layer = create_embedding_layer(embedding_size, embedding_dim)
+
+        input_shape = (2, 6)
+        output_shape = layer.compute_output_shape(input_shape)
+        self.assertEqual(output_shape, input_shape + (embedding_dim,))
+
+        ids = [0, 1, 3, 8, 3, 2, 3]
+        values = layer.call(ids)
+        values = values.numpy()
+        for index, idx in enumerate(ids):
+            correct_value = np.array([idx] * embedding_dim, dtype=np.float32)
+            self.assertTrue((values[index] == correct_value).all())
+
+        model = tf.keras.models.Sequential([layer])
+        outputs = model.call(np.array([ids, ids]))
+        values = outputs.numpy()[1]
+        for index, idx in enumerate(ids):
+            correct_value = np.array([idx] * embedding_dim, dtype=np.float32)
+            self.assertTrue((values[index] == correct_value).all())
+
+    def test_embedding_layer_gradient(self):
+        embedding_dim = 8
+        embedding_size = 16
+        layer = create_embedding_layer(embedding_size, embedding_dim)
+        inputs_list = [
+            tf.keras.backend.constant([[0, 1, 3], [1, 2, 0]], dtype=tf.int64),
+            tf.keras.backend.constant(
+                [[0, 10, 3], [11, 2, 1]], dtype=tf.int64
+            ),
+        ]
+        multiply_values = np.ndarray(
+            shape=(6, embedding_dim), dtype=np.float32
+        )
+        for i in range(6):
+            multiply_values[i].fill(i)
+        multiply_tensor = tf.convert_to_tensor(multiply_values)
+        multiply_tensor = tf.reshape(multiply_tensor, [2, 3, 8])
+
+        for inputs in inputs_list:
+            with tf.GradientTape() as tape:
+                layer.set_tape(tape)
+                output = layer.call(inputs)
+                output = output * multiply_tensor
+            bet = layer.bet_ids_pair[0][0]
+            grads = tape.gradient(output, bet)
+            layer.reset()
+            self.assertTrue((grads.values.numpy() == multiply_values).all())
+
+
 if __name__ == "__main__":
     unittest.main()