Repository: sql-machine-learning/elasticdl

Issue:
Disable tf.function when EdlEmbedding is used



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
@@ -33,15 +33,23 @@ def name(self):
     def get_key(name_list):
         return "-".join(map(str, name_list))
 
+    def lookup_embedding(self, unique_ids):
+        batch_embedding = self.worker.embedding_lookup(
+            unique_ids, self._name, self.embedding_initializer
+        )
+        return batch_embedding
+
     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.py_function(
+            self.lookup_embedding, inp=[unique_ids], Tout=tf.float32
         )
-        batch_embedding_tensor = tf.convert_to_tensor(batch_embedding)
         if self.tape:
+            # tape.watch works with eager mode only
+            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))
         outputs = tf.gather(batch_embedding_tensor, idx)
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
@@ -255,8 +255,21 @@ def get_trainable_items(self):
                 bets.extend([i for (i, _) in layer.bet_ids_pair])
         return self._model.trainable_variables + bets
 
-    @tf.function
     def training_process(self, features, labels):
+        """
+        training for models with elasticdl.layers.embedding does not
+        support tf.function decorator
+        """
+        if self._embedding_layers:
+            return self.training_process_eagerly(features, labels)
+        else:
+            return self.training_process_with_acceleration(features, labels)
+
+    @tf.function
+    def training_process_with_acceleration(self, features, labels):
+        return self.training_process_eagerly(features, labels)
+
+    def training_process_eagerly(self, features, labels):
         with tf.GradientTape() as tape:
             self._set_tape_for_embedding(tape)
             outputs = self._model.call(features, training=True)


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
@@ -77,6 +77,11 @@ def create_embedding_layer(embedding_size, embedding_dim):
     return layer
 
 
+@tf.function
+def layer_call(layer, inputs):
+    return layer.call(inputs)
+
+
 class EmbeddingLayerTest(unittest.TestCase):
     def test_embedding_layer(self):
         embedding_dim = 8
@@ -94,6 +99,11 @@ def test_embedding_layer(self):
             correct_value = np.array([idx] * embedding_dim, dtype=np.float32)
             self.assertTrue((values[index] == correct_value).all())
 
+        results = layer_call(layer, ids)
+        results = results.numpy()
+        for index, idx in enumerate(ids):
+            self.assertTrue((values[index] == results[index]).all())
+
         model = tf.keras.models.Sequential([layer])
         outputs = model.call(np.array([ids, ids]))
         values = outputs.numpy()[1]
@@ -129,6 +139,11 @@ def test_embedding_layer_gradient(self):
             layer.reset()
             self.assertTrue((grads.values.numpy() == multiply_values).all())
 
+        for inputs in inputs_list:
+            with tf.GradientTape() as tape:
+                layer.set_tape(tape)
+                self.assertRaises(RuntimeError, layer_call, layer, inputs)
+
 
 if __name__ == "__main__":
     unittest.main()