Repository: sql-machine-learning/elasticdl

Issue:
Replace `model.trainable_variables` with non-embedding trainable variables
Currently, `model.trainable_variables` are used in worker, e.g. `get_model` function 

https://github.com/sql-machine-learning/elasticdl/blob/7ebcea845bbb898d3278af4bbfff33b51bdaeb23/elasticdl/python/worker/worker.py#L135-L147

and `report_variable` function

https://github.com/sql-machine-learning/elasticdl/blob/7ebcea845bbb898d3278af4bbfff33b51bdaeb23/elasticdl/python/worker/worker.py#L208-L215

However, these functions should not use trainable variables from ElasticDL embedding layers. This is because we save variables of ElasticDL embedding layer in Redis. 

Now ElasticDL embedding layer does not have trainable variables. But it will have trainable variables in order to support `tf.function`. 

So we should replace `model.trainable_variables` with non-embedding trainable variables. We can add a function to get these variables and save them in worker.


Discussion:

Resolving patch:
diff --git a/elasticdl/python/common/model_helper.py b/elasticdl/python/common/model_helper.py
--- a/elasticdl/python/common/model_helper.py
+++ b/elasticdl/python/common/model_helper.py
@@ -152,3 +152,15 @@ def find_layer(model, layer_class):
             # search in nested layers
             layers += find_layer(layer, layer_class)
     return layers
+
+
+def get_non_embedding_trainable_vars(model, embedding_layers):
+    """
+    Get trainable variables which are not from ElasticDL embedding layers.
+    """
+    embedding_items = []
+    for layer in embedding_layers:
+        embedding_items.extend([bet for bet, _ in layer.bet_ids_pair])
+    return [
+        var for var in model.trainable_variables if var not in embedding_items
+    ]
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
@@ -6,7 +6,11 @@
 from elasticdl.proto import elasticdl_pb2, elasticdl_pb2_grpc
 from elasticdl.python.common.constants import JobType, Mode
 from elasticdl.python.common.log_util import default_logger as logger
-from elasticdl.python.common.model_helper import find_layer, get_model_spec
+from elasticdl.python.common.model_helper import (
+    find_layer,
+    get_model_spec,
+    get_non_embedding_trainable_vars,
+)
 from elasticdl.python.common.ndarray import (
     ndarray_to_tensor,
     tensor_to_ndarray,
@@ -73,7 +77,7 @@ def __init__(
         self._job_type = job_type
         self._minibatch_size = minibatch_size
         (
-            self._model,
+            model_inst,
             self._dataset_fn,
             self._loss,
             self._opt_fn,
@@ -89,8 +93,7 @@ def __init__(
             model_params=model_params,
             prediction_outputs_processor=prediction_outputs_processor,
         )
-        self._init_embedding_layer()
-        self._var_created = self._model.built
+        self.set_model(model_inst)
 
         if channel is None:
             self._stub = None
@@ -104,6 +107,19 @@ def __init__(
         )
         self._get_model_steps = get_model_steps
 
+    # TODO: Multiple tests are currently using this function to initialize
+    # self._model, where the initialization should be done via constructor.
+    def set_model(self, model_inst):
+        """Set model instance to worker."""
+        self._model = model_inst
+        self._init_embedding_layer()
+        self._var_created = self._model.built
+        self._non_embed_vars = []
+        if self._var_created:
+            self._non_embed_vars = get_non_embedding_trainable_vars(
+                self._model, self._embedding_layers
+            )
+
     def _init_embedding_layer(self):
         """
         Init elasticdl.layers.embedding layer list and assign worker to them
@@ -141,7 +157,7 @@ def get_model(self, version, method):
         req.method = method
         model = self._stub.GetModel(req)
 
-        for var in self._model.trainable_variables:
+        for var in self._non_embed_vars:
             # Assumes all trainable variables exist in model.param.
             var.assign(tensor_to_ndarray(model.param[var.name]))
         self._model_version = model.version
@@ -210,7 +226,7 @@ def report_variable(self):
         report variable to ps.
         """
         req = elasticdl_pb2.ReportVariableRequest()
-        for v in self._model.trainable_variables:
+        for v in self._non_embed_vars:
             req.variable[v.name].CopyFrom(ndarray_to_tensor(v.numpy()))
         self._stub.ReportVariable(req)
 
@@ -219,10 +235,10 @@ def report_gradient(self, grads):
         report gradient to ps, return (accepted, model_version) from rpc call.
         """
         req = elasticdl_pb2.ReportGradientRequest()
-        origin_vars = self._model.trainable_variables
-        origin_var_n = len(origin_vars)
-        # should keep the same order as self.get_trainable_items()
-        for g, v in zip(grads[:origin_var_n], origin_vars):
+        non_embed_vars_n = len(self._non_embed_vars)
+        # The first `non_embed_vars_n` items in `grads` are gradients for
+        # `self._non_embed_vars`
+        for g, v in zip(grads[:non_embed_vars_n], self._non_embed_vars):
             if isinstance(g, tf.IndexedSlices):
                 req.gradient[v.name].CopyFrom(
                     ndarray_to_tensor(
@@ -233,10 +249,14 @@ def report_gradient(self, grads):
                 req.gradient[v.name].CopyFrom(ndarray_to_tensor(g.numpy()))
 
         # Accumulate gradients of ElasticDL embedding layer
-        # should keep the same order as self.get_trainable_items()
         if self._embedding_layers:
-            edl_embedding_grads = grads[origin_var_n:]
+            # The `edl_embedding_grads` are gradients for bets in
+            # `self._embedding_layers`
+            edl_embedding_grads = grads[non_embed_vars_n:]
 
+            # Check that the number of bet equal to the number of gradients.
+            # Please note that every embedding layer may have more than one
+            # `bet_id_pair`.
             bet_number = 0
             for layer in self._embedding_layers:
                 bet_number += len(layer.bet_ids_pair)
@@ -306,6 +326,9 @@ def report_prediction_outputs(self, predictions):
     def _create_variable_and_report(self, features):
         # Use model.call to create variables, then report to ps
         _ = self._model.call(features)
+        self._non_embed_vars = get_non_embedding_trainable_vars(
+            self._model, self._embedding_layers
+        )
         self.report_variable()
         self._var_created = True
 
@@ -319,7 +342,7 @@ def get_trainable_items(self):
         if self._embedding_layers:
             for layer in self._embedding_layers:
                 bets.extend([i for (i, _) in layer.bet_ids_pair])
-        return self._model.trainable_variables + bets
+        return self._non_embed_vars + bets
 
     def training_process(self, features, labels):
         """


Test patch:
diff --git a/elasticdl/python/tests/indices_slices_gradient_test.py b/elasticdl/python/tests/indices_slices_gradient_test.py
--- a/elasticdl/python/tests/indices_slices_gradient_test.py
+++ b/elasticdl/python/tests/indices_slices_gradient_test.py
@@ -48,7 +48,7 @@ def _create_master_and_worker(self):
             model_def="test_module.custom_model",
             channel=None,
         )
-        worker._model = model_inst
+        worker.set_model(model_inst)
         worker._stub = InProcessMaster(master)
 
         return master, worker
diff --git a/elasticdl/python/tests/optimizer_wrapper_test.py b/elasticdl/python/tests/optimizer_wrapper_test.py
--- a/elasticdl/python/tests/optimizer_wrapper_test.py
+++ b/elasticdl/python/tests/optimizer_wrapper_test.py
@@ -21,6 +21,7 @@
 from elasticdl.python.common.model_helper import (
     find_layer,
     get_module_file_path,
+    get_non_embedding_trainable_vars,
     load_module,
 )
 from elasticdl.python.elasticdl.layers.embedding import Embedding
@@ -113,54 +114,54 @@ def lookup_func(ids, layer_name, initializer, output_dim):
         layer.set_lookup_func(lookup_func)
 
     # training process
-    for features, labels in zip(X, Y):
+    for train_iter, (features, labels) in enumerate(zip(X, Y)):
         with tf.GradientTape() as tape:
             for layer in embed_layers:
                 layer.set_tape(tape)
             outputs = model.call(features)
             loss = loss_fn(outputs, labels)
 
-        # TODO: calculate train_vars_embed and train_vars_other can be a
-        # reusable function
-        train_vars_embed = []
-        train_vars_other = []
-        for layer in model.layers:
-            if isinstance(layer, Embedding):
-                for bet, ids in layer.bet_ids_pair:
-                    train_vars_embed.append((bet, layer.name, ids))
-            else:
-                vars = layer.trainable_variables
-                train_vars_other.extend(vars)
+        # Need to get non-embedding variables inside for loop because model
+        # creates variables after the first time `model.call` is called
+        if not train_iter:
+            non_embed_vars = get_non_embedding_trainable_vars(
+                model, embed_layers
+            )
+        embed_items = []
+        for layer in embed_layers:
+            embed_items.extend(
+                [(bet, layer.name, ids) for bet, ids in layer.bet_ids_pair]
+            )
 
         grads = tape.gradient(
-            loss, train_vars_other + [var for var, _, _ in train_vars_embed]
+            loss, non_embed_vars + [var for var, _, _ in embed_items]
         )
 
         # TODO: do not need to merge gradient from the same embedding layer
         # after `optimizer_wrapper` support grads_and_vars with duplicated
         # layer name
-        train_vars_other_len = len(train_vars_other)
-        grads_new = grads[:train_vars_other_len]
-        grads_embed_dict = {}
+        non_embed_vars_n = len(non_embed_vars)
+        non_embed_grads = grads[:non_embed_vars_n]
+        embed_grads_dict = {}
         for (_, layer_name, ids), grad in zip(
-            train_vars_embed, grads[train_vars_other_len:]
+            embed_items, grads[non_embed_vars_n:]
         ):
-            if layer_name in grads_embed_dict:
-                grads_merged = grads_embed_dict[layer_name]
-                grads_embed_dict[layer_name] = tf.IndexedSlices(
-                    tf.concat([grads_merged.values, grad.values], axis=0),
-                    tf.concat([grads_merged.indices, ids], axis=0),
+            if layer_name in embed_grads_dict:
+                merged_grads = embed_grads_dict[layer_name]
+                embed_grads_dict[layer_name] = tf.IndexedSlices(
+                    tf.concat([merged_grads.values, grad.values], axis=0),
+                    tf.concat([merged_grads.indices, ids], axis=0),
                 )
             else:
-                grads_embed_dict[layer_name] = tf.IndexedSlices(
+                embed_grads_dict[layer_name] = tf.IndexedSlices(
                     grad.values, ids
                 )
 
         optimizer.apply_gradients(
-            list(zip(grads_new, train_vars_other))
+            list(zip(non_embed_grads, non_embed_vars))
             + [
                 (grad, layer_name)
-                for layer_name, grad in grads_embed_dict.items()
+                for layer_name, grad in embed_grads_dict.items()
             ]
         )
 
diff --git a/elasticdl/python/tests/report_gradients_of_bet_test.py b/elasticdl/python/tests/report_gradients_of_bet_test.py
--- a/elasticdl/python/tests/report_gradients_of_bet_test.py
+++ b/elasticdl/python/tests/report_gradients_of_bet_test.py
@@ -86,7 +86,7 @@ def _create_master_and_worker(self):
             model_def="test_module.custom_model",
             channel=None,
         )
-        worker._model = model_inst
+        worker.set_model(model_inst)
         worker._stub = InProcessMaster(master)
 
         return master, worker