Repository: sql-machine-learning/elasticdl

Issue:
Do not save ElasticDL embedding's variables in the master
In current code, ElasticDL embedding layer create trainable variables and use them only for training under graph mode (i.e. under tf.function decorator). These variables are used for training and thus should only be created in workers. 

However, now ElasticDL embedding layer creates these variables in `__init__` function. The master (or PS) creates these variables when it creates a model instance, and then save these variables in `self._model`. 

I propose that ElasticDL embedding layer should not create these variables in `__init__` function. Instead, creating them before the training process only in worker is enough.


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
@@ -64,16 +64,21 @@ def __init__(
         # `None` dimension. This is because they have different shapes in
         # different iterations.
         # `tf.Variable` requires initial value if shape has `None` dimension.
+        self._embedding_and_ids_graph = []
+
+    def _init_for_graph_mode(self):
         self._embedding_and_ids_graph = [
             EmbeddingAndIds(
                 batch_embedding=tf.Variable(
-                    initial_value=tf.zeros((1, self.output_dim)),
+                    # In some cases, `tf.Variable` requires that initial value
+                    # is callable.
+                    initial_value=lambda: tf.zeros((1, self.output_dim)),
                     shape=tf.TensorShape((None, self.output_dim)),
                     dtype=tf.float32,
                     trainable=True,
                 ),
                 batch_ids=tf.Variable(
-                    initial_value=tf.zeros((1, 1), dtype=tf.int64),
+                    initial_value=lambda: tf.zeros((1, 1), dtype=tf.int64),
                     shape=tf.TensorShape(None),
                     dtype=tf.int64,
                     trainable=False,
@@ -186,6 +191,13 @@ def _record_gradients(self, batch_embedding, ids):
         return batch_embedding
 
     def call(self, input):
+        if (
+            self.tape
+            and not tf.executing_eagerly()
+            and not self._embedding_and_ids_graph
+        ):
+            self._init_for_graph_mode()
+
         input = tf.cast(input, tf.int64)
         if isinstance(input, tf.SparseTensor):
             return self._sparse_input_call(input)


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
@@ -88,11 +88,6 @@ def create_embedding_layer(
     return layer
 
 
-@tf.function
-def module_call(module, inputs):
-    return module.call(inputs)
-
-
 def get_correct_values_for_sparse_test(indices, values, dense_shape, combiner):
     results = []
     for i in range(dense_shape[0]):
@@ -116,10 +111,10 @@ def get_correct_values_for_sparse_test(indices, values, dense_shape, combiner):
 
 class EmbeddingLayerTest(unittest.TestCase):
     def _run_forward_pass_and_compare(
-        self, call_fns, correct_values, output_dim
+        self, call_fns, inputs, correct_values, output_dim
     ):
         for call_fn in call_fns:
-            values = call_fn()
+            values = call_fn(inputs)
             values = values.numpy().reshape(-1, output_dim)
             for v, correct_v in zip(values, correct_values):
                 self.assertTrue(np.isclose(v, correct_v).all())
@@ -134,12 +129,12 @@ def test_embedding_layer(self):
         self.assertEqual(output_shape, input_shape + (output_dim,))
 
         ids = [0, 1, 3, 8, 3, 2, 3]
-        call_fns = [lambda: layer.call(ids), lambda: module_call(layer, ids)]
+        call_fns = [layer.call, tf.function(layer.call)]
         correct_values = np.array(
             [np.array([idx] * output_dim, dtype=np.float32) for idx in ids]
         )
         self._run_forward_pass_and_compare(
-            call_fns, correct_values, output_dim
+            call_fns, ids, correct_values, output_dim
         )
 
         # Keras model without/with input_layer
@@ -152,12 +147,9 @@ def test_embedding_layer(self):
         models = [model_without_input_layer, model_with_input_layer]
         for model in models:
             inputs = tf.constant([ids])
-            call_fns = [
-                lambda: model.call(inputs),
-                lambda: module_call(model, inputs),
-            ]
+            call_fns = [model.call, tf.function(model.call)]
             self._run_forward_pass_and_compare(
-                call_fns, correct_values, output_dim
+                call_fns, inputs, correct_values, output_dim
             )
 
     def test_embedding_layer_with_input_length(self):
@@ -169,7 +161,7 @@ def test_embedding_layer_with_input_length(self):
         )
         ids = [[0, 1, 3, 8], [5, 3, 2, 3]]
         flatten_ids = ids[0] + ids[1]
-        call_fns = [lambda: layer.call(ids), lambda: module_call(layer, ids)]
+        call_fns = [layer.call, tf.function(layer.call)]
         correct_values = np.array(
             [
                 np.array([idx] * output_dim, dtype=np.float32)
@@ -177,7 +169,7 @@ def test_embedding_layer_with_input_length(self):
             ]
         )
         self._run_forward_pass_and_compare(
-            call_fns, correct_values, output_dim
+            call_fns, ids, correct_values, output_dim
         )
 
     def test_embedding_layer_with_sparse_input(self):
@@ -195,10 +187,7 @@ def test_embedding_layer_with_sparse_input(self):
             inputs = tf.SparseTensor(
                 indices=indices, values=values, dense_shape=dense_shape
             )
-            call_fns = [
-                lambda: layer.call(inputs),
-                lambda: module_call(layer, inputs),
-            ]
+            call_fns = [layer.call, tf.function(layer.call)]
             correct_value_single_line = get_correct_values_for_sparse_test(
                 indices, values, dense_shape, combiner
             )
@@ -209,7 +198,7 @@ def test_embedding_layer_with_sparse_input(self):
                 ]
             )
             self._run_forward_pass_and_compare(
-                call_fns, correct_values, output_dim
+                call_fns, inputs, correct_values, output_dim
             )
 
     def test_embedding_layer_with_mask_zero(self):
@@ -253,14 +242,11 @@ def test_embedding_layer_gradient(self):
         multiply_tensor = tf.reshape(multiply_tensor, [2, 3, 8])
 
         for inputs in inputs_list:
-            call_fns = [
-                lambda: layer.call(inputs),
-                lambda: module_call(layer, inputs),
-            ]
+            call_fns = [layer.call, tf.function(layer.call)]
             for call_fn in call_fns:
                 with tf.GradientTape() as tape:
                     layer.set_tape(tape)
-                    output = module_call(layer, inputs)
+                    output = call_fn(inputs)
                     output = output * multiply_tensor
                 batch_embedding = layer.embedding_and_ids[0].batch_embedding
                 grads = tape.gradient(output, batch_embedding)
@@ -329,14 +315,11 @@ def test_embedding_layer_gradient_with_sparse_inputs(self):
             layer = create_embedding_layer(
                 embedding_size, output_dim, combiner=combiner
             )
-            call_fns = [
-                lambda: layer.call(inputs),
-                lambda: module_call(layer, inputs),
-            ]
+            call_fns = [tf.function(layer.call), layer.call]
             for call_fn in call_fns:
                 with tf.GradientTape() as tape:
                     layer.set_tape(tape)
-                    output = module_call(layer, inputs)
+                    output = call_fn(inputs)
                     output = output * multiply_tensor
                 batch_embedding = layer.embedding_and_ids[0].batch_embedding
                 grads = tape.gradient(output, batch_embedding)