Repository: sql-machine-learning/elasticdl

Issue:
Worker does backward-pass and report gradients to master
This issue is about [worker computes embedding-vectors gradient](https://github.com/wangkuiyi/elasticdl/blob/develop/elasticdl/doc/distributed_embedding_layer_design.md#worker-computes-embedding-vectors-gradient) and [worker reports embedding vectors gradient to master](https://github.com/wangkuiyi/elasticdl/blob/develop/elasticdl/doc/distributed_embedding_layer_design.md#worker-reports-embedding-vectors-gradient-to-master) in design doc. 
Worker does backward-pass, extracts gradients from `elasticdl.Embedding` layer and reports them to master.


Discussion:

Resolving patch:
diff --git a/elasticdl/python/common/tensor_helper.py b/elasticdl/python/common/tensor_helper.py
new file mode 100644
--- /dev/null
+++ b/elasticdl/python/common/tensor_helper.py
@@ -0,0 +1,8 @@
+import tensorflow as tf
+
+
+def merge_indexed_slices(*args):
+    return tf.IndexedSlices(
+        tf.concat([i.values for i in args], axis=0),
+        tf.concat([i.indices for i in args], axis=0),
+    )
diff --git a/elasticdl/python/master/servicer.py b/elasticdl/python/master/servicer.py
--- a/elasticdl/python/master/servicer.py
+++ b/elasticdl/python/master/servicer.py
@@ -12,6 +12,7 @@
     ndarray_to_tensor,
     tensor_to_ndarray,
 )
+from elasticdl.python.common.tensor_helper import merge_indexed_slices
 from elasticdl.python.master.checkpoint_service import CheckpointService
 
 
@@ -35,6 +36,7 @@ def __init__(
         self._task_d = task_d
         self._lock = threading.Lock()
         self._gradient_sum = {}
+        self._edl_embedding_gradients = {}
         self._gradient_sum_indexed = {}
         self._grad_to_wait = grads_to_wait
         self._grad_n = 0
@@ -150,6 +152,7 @@ def _update_model(self):
         self._update_model_version()
         self._gradient_sum.clear()
         self._gradient_sum_indexed.clear()
+        self._edl_embedding_gradients.clear()
         self._grad_n = 0
 
     def get_model_version(self):
@@ -248,9 +251,29 @@ def ReportGradient(self, request, _):
             # Do sanity check before accumulating gradients.
             for k, v in request.gradient.items():
                 if k not in self._model:
-                    raise ValueError(
-                        "Gradient key: %s is not part of model", k
-                    )
+                    if v.indices:
+                        # IndexedSlices are gradients of Embeddings Table
+                        # of EdlEmbedding layer
+                        # TODO: Updating Embedding Table of EdlEmbedding
+                        # layer will be implemented
+                        arr = tensor_to_ndarray(v)
+                        if k in self._edl_embedding_gradients:
+                            self._edl_embedding_gradients[
+                                k
+                            ] = merge_indexed_slices(
+                                self._edl_embedding_gradients[k], arr
+                            )
+                        else:
+                            self._edl_embedding_gradients[k] = arr
+                        logger.warning(
+                            "Update Embedding Table of EdlEmbedding layer "
+                            "is not implemented."
+                        )
+                        continue
+                    else:
+                        raise ValueError(
+                            "Gradient key: %s is not part of model", k
+                        )
                 arr = tensor_to_ndarray(v)
                 if isinstance(arr, tf.IndexedSlices):
                     if arr.values.shape[1] != self._model[k].numpy().shape[1]:
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
@@ -56,6 +56,8 @@ def __init__(
         self._model = load_model_from_module(
             model_def, model_module, model_params
         )
+        # TODO: call find_layer and get EdlEmbedding layers
+        self._embedding_layers = []
         self._var_created = self._model.built
         self._dataset_fn = model_module[dataset_fn]
         self._opt_fn = model_module[optimizer]
@@ -133,7 +135,10 @@ def report_gradient(self, grads):
         report gradient to ps, return (accepted, model_version) from rpc call.
         """
         req = elasticdl_pb2.ReportGradientRequest()
-        for g, v in zip(grads, self._model.trainable_variables):
+        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):
             if isinstance(g, tf.IndexedSlices):
                 req.gradient[v.name].CopyFrom(
                     ndarray_to_tensor(
@@ -142,6 +147,43 @@ def report_gradient(self, grads):
                 )
             else:
                 req.gradient[v.name].CopyFrom(ndarray_to_tensor(g.numpy()))
+
+        # should keep the same order as self.get_trainable_items()
+        if self._embedding_layers:
+            grads_edlembedding = grads[origin_var_n:]
+
+            bet_number = 0
+            for layer in self._embedding_layers:
+                bet_number += len(layer.bet_ids_pair)
+            if len(grads_edlembedding) != bet_number:
+                raise ValueError(
+                    "EdlEmbedding related gradient number %d does not match"
+                    "the number of EdlEmbedding output tensor %d."
+                    % (len(grads_edlembedding), bet_number)
+                )
+
+            it = 0
+            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:
+                    if g_values is not None:
+                        g_values = tf.concat(
+                            [g_values, grads_edlembedding[it]], axis=0
+                        )
+                        g_indices = tf.concat([g_indices, ids], axis=0)
+                    else:
+                        g_values = grads_edlembedding[it]
+                        g_indices = ids
+                    it += 1
+
+                req.gradient[layer.name].CopyFrom(
+                    ndarray_to_tensor(
+                        g_values.numpy(), tuple(g_indices.numpy())
+                    )
+                )
+
         req.model_version = self._model_version
         res = self._stub.ReportGradient(req)
         return res.accepted, res.model_version
@@ -181,6 +223,18 @@ def _create_variable_and_report(self, features):
         self.report_variable()
         self._var_created = True
 
+    def get_trainable_items(self):
+        """
+        return all trainable variables list, including batch embedding
+        tensor (BET) if exists. take care to keep the same order as in
+        self.report_gradient()
+        """
+        bets = []
+        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
+
     @tf.function
     def training_process(self, features, labels):
         with tf.GradientTape() as tape:
@@ -189,7 +243,7 @@ def training_process(self, features, labels):
             # Add regularization loss if any
             if self._model.losses:
                 loss += tf.math.add_n(self._model.losses)
-        grads = tape.gradient(loss, self._model.trainable_variables)
+        grads = tape.gradient(loss, self.get_trainable_items())
         return loss, grads
 
     @tf.function


Test patch:
diff --git a/elasticdl/python/tests/report_gradients_of_bet_test.py b/elasticdl/python/tests/report_gradients_of_bet_test.py
new file mode 100644
--- /dev/null
+++ b/elasticdl/python/tests/report_gradients_of_bet_test.py
@@ -0,0 +1,193 @@
+import os
+import unittest
+
+import numpy as np
+import tensorflow as tf
+
+from elasticdl.python.common.constants import JobType
+from elasticdl.python.common.model_helper import get_model_file, load_module
+from elasticdl.python.master.servicer import MasterServicer
+from elasticdl.python.tests.in_process_master import InProcessMaster
+from elasticdl.python.worker.worker import Worker
+
+_model_file = get_model_file(
+    os.path.dirname(os.path.realpath(__file__)), "test_module.custom_model"
+)
+m = load_module(_model_file).__dict__
+
+
+def custom_model():
+    embedding_weight = np.array(range(1, 13), dtype=np.float32).reshape(4, 3)
+    dense_weight = np.array(range(13, 19), dtype=np.float32).reshape(-1, 1)
+
+    inputs = tf.keras.Input(shape=(2,), name="image")
+    x = tf.keras.layers.Embedding(
+        input_dim=4, output_dim=3, input_length=2, weights=[embedding_weight]
+    )(inputs)
+    x = tf.keras.layers.Flatten()(x)
+    outputs = tf.keras.layers.Dense(1, use_bias=False, weights=[dense_weight])(
+        x
+    )
+    return tf.keras.Model(inputs=inputs, outputs=outputs)
+
+
+class MockEdlEmbedding:
+    def __init__(self, name):
+        self._name = name
+        self.bet_ids_pair = []
+
+    @property
+    def name(self):
+        return self._name
+
+    def add_bet_and_ids(self, bet, ids):
+        self.append((bet, ids))
+
+
+class ReportBETGradientTest(unittest.TestCase):
+    def _create_master_and_worker(self):
+        model_inst = custom_model()
+        master = MasterServicer(
+            2,
+            2,
+            tf.optimizers.SGD(0.1),
+            None,
+            init_var=model_inst.trainable_variables,
+            checkpoint_filename_for_init=None,
+            checkpoint_service=None,
+            evaluation_service=None,
+        )
+        worker = Worker(
+            1,
+            JobType.TRAINING_ONLY,
+            2,
+            _model_file,
+            model_def="test_module.custom_model",
+            channel=None,
+        )
+        worker._model = model_inst
+        worker._stub = InProcessMaster(master)
+
+        return master, worker
+
+    def test_report_bet_gradients(self):
+        master, worker = self._create_master_and_worker()
+
+        layer1 = MockEdlEmbedding("test_edlembedding_1")
+        layer1.bet_ids_pair = [
+            (None, tf.constant([1, 2])),
+            (None, tf.constant([2, 3])),
+        ]
+
+        layer2 = MockEdlEmbedding("test_edlembedding_2")
+        layer2.bet_ids_pair = [
+            (None, tf.constant([3, 1])),
+            (None, tf.constant([3, 4])),
+        ]
+
+        edlembed_grads = [
+            tf.reshape(tf.range(1, 7, dtype=tf.float32), (2, 3)),
+            tf.reshape(tf.range(7, 13, dtype=tf.float32), (2, 3)),
+            tf.reshape(tf.range(13, 19, dtype=tf.float32), (2, 3)),
+            tf.reshape(tf.range(19, 25, dtype=tf.float32), (2, 3)),
+        ]
+
+        worker._embedding_layers = [layer1, layer2]
+
+        values1 = tf.convert_to_tensor(
+            np.array([[1, 2, 3], [10, 11, 12]], dtype=np.float32)
+        )
+        indices1 = tf.convert_to_tensor([0, 3], dtype=tf.int64)
+
+        grads1 = [
+            tf.IndexedSlices(values1, indices1),
+            tf.reshape(
+                tf.convert_to_tensor(range(1, 7), dtype=tf.float32),
+                shape=(6, 1),
+            ),
+        ]
+        grads1.extend(edlembed_grads)
+
+        worker._model_version = 0
+        worker.report_gradient(grads1)
+
+        expected_edlembedding_grads = {
+            layer1.name: tf.IndexedSlices(
+                tf.concat(edlembed_grads[:2], axis=0),
+                tf.concat(
+                    [layer1.bet_ids_pair[0][1], layer1.bet_ids_pair[1][1]],
+                    axis=0,
+                ),
+            ),
+            layer2.name: tf.IndexedSlices(
+                tf.concat(edlembed_grads[2:], axis=0),
+                tf.concat(
+                    [layer2.bet_ids_pair[0][1], layer2.bet_ids_pair[1][1]],
+                    axis=0,
+                ),
+            ),
+        }
+
+        result = master._edl_embedding_gradients
+        for name, grads in expected_edlembedding_grads.items():
+            self.assertTrue(name in result)
+            self.assertTrue(grads.indices.shape == result[name].indices.shape)
+            self.assertTrue(grads.values.shape == result[name].values.shape)
+            self.assertTrue(
+                np.all(tf.equal(grads.indices, result[name].indices))
+            )
+            self.assertTrue(
+                np.all(grads.values - result[name].values < 0.0001)
+            )
+
+        # make sure other gradients are calculated correctly
+        values2 = tf.convert_to_tensor(
+            np.array([[7, 8, 9], [4, 5, 6]], dtype=np.float32)
+        )
+        indices2 = tf.convert_to_tensor([2, 0], dtype=tf.int64)
+        grads2 = [
+            tf.IndexedSlices(values2, indices2),
+            tf.reshape(
+                tf.convert_to_tensor(range(13, 19), dtype=tf.float32),
+                shape=(6, 1),
+            ),
+        ]
+        grads2.extend(edlembed_grads)
+        worker.report_gradient(grads2)
+
+        expected_weights = []
+        expected_weights.append(
+            np.array(
+                [
+                    [0.5, 1.3, 2.1],
+                    [4.0, 5.0, 6.0],
+                    [6.3, 7.2, 8.1],
+                    [9.0, 9.9, 10.8],
+                ]
+            )
+        )
+        expected_weights.append(
+            np.reshape(
+                np.array([[12.3, 13.2, 14.1, 15.0, 15.9, 16.8]]), (6, 1)
+            )
+        )
+
+        for i, j in zip(master._model.values(), expected_weights):
+            self.assertTrue(np.all(i.numpy() - j < 0.0001))
+
+    def test_get_trainable_variable(self):
+        master, worker = self._create_master_and_worker()
+        layer = MockEdlEmbedding("test")
+        layer.bet_ids_pair = [
+            (tf.Variable([1, 2, 3], name="test_bet"), [1, 2, 3])
+        ]
+        worker._embedding_layers = [layer]
+        train_vars = worker.get_trainable_items()
+        print(train_vars)
+        self.assertTrue("embedding" in train_vars[0].name)
+        self.assertTrue("dense" in train_vars[1].name)
+        self.assertTrue("test_bet" in train_vars[2].name)
+
+
+if __name__ == "__main__":
+    unittest.main()