Repository: sql-machine-learning/elasticdl

Issue:
Compute Embedding Layer's Gradients id in Embedding Table
The embedding layer's gradient from the tape uses ids of *BET*. We need to transfer these ids to the corresponding embedding table id.


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
@@ -88,7 +88,10 @@ def get_key(name_list):
 
     def lookup_embedding(self, unique_ids):
         batch_embedding = self.worker.lookup_embedding(
-            unique_ids, self._name, self.embedding_initializer, self.output_dim
+            unique_ids.numpy(),
+            self._name,
+            self.embedding_initializer,
+            self.output_dim,
         )
         return batch_embedding
 
@@ -110,7 +113,7 @@ def call(self, input):
             if not tf.executing_eagerly():
                 raise RuntimeError("tape.watch only works with eager mode")
             self.tape.watch(batch_embedding_tensor)
-            self.bet_ids_pair.append((batch_embedding_tensor, unique_ids))
+            self.bet_ids_pair.append((batch_embedding_tensor, flat_ids))
         outputs = tf.gather(batch_embedding_tensor, idx)
         # tf.reshape does not support shape with None. Replace None with -1.
         if ids.get_shape().rank == 2:
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
@@ -217,6 +217,7 @@ def report_gradient(self, grads):
             else:
                 req.gradient[v.name].CopyFrom(ndarray_to_tensor(g.numpy()))
 
+        # deal with gradients of ElasticDL embedding layer
         # should keep the same order as self.get_trainable_items()
         if self._embedding_layers:
             grads_edlembedding = grads[origin_var_n:]
@@ -235,17 +236,18 @@ def report_gradient(self, grads):
             for layer in self._embedding_layers:
                 g_values = None
                 g_indices = None
-                # TODO: translate ids in gradients of bet
                 for bet, ids in layer.bet_ids_pair:
+                    grad = grads_edlembedding[it]
+                    it += 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, grads_edlembedding[it]], axis=0
-                        )
+                        g_values = tf.concat([g_values, grad], axis=0)
                         g_indices = tf.concat([g_indices, ids], axis=0)
                     else:
-                        g_values = grads_edlembedding[it]
+                        g_values = grad
                         g_indices = ids
-                    it += 1
 
                 req.gradient[layer.name].CopyFrom(
                     ndarray_to_tensor(


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
@@ -237,8 +237,14 @@ def test_embedding_layer_gradient(self):
                 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())
+            self.assertTrue(
+                (
+                    layer.bet_ids_pair[0][1].numpy()
+                    == inputs.numpy().reshape(-1)
+                ).all()
+            )
+            layer.reset()
 
         for inputs in inputs_list:
             with tf.GradientTape() as tape:
@@ -288,10 +294,17 @@ def test_embedding_layer_gradient_with_sparse_inputs(self):
             bet = layer.bet_ids_pair[0][0]
             grads = tape.gradient(output, bet)
             grads = grads.numpy()
-            layer.reset()
             place = 8 if combiner == "sum" else 5
             for n, v in enumerate(correct_grads[combiner]):
                 self.assertAlmostEqual(grads[n][0], v, place)
+            self.assertTrue(
+                (
+                    layer.bet_ids_pair[0][1].numpy()
+                    - np.array([1, 3, 2, 0, 6])
+                    < 0.00001
+                ).all()
+            )
+            layer.reset()
 
 
 if __name__ == "__main__":