Repository: sql-machine-learning/elasticdl

Issue:
Support Training model with Keras Embedding layer
In our implementation now, the gradients are supposed to be in the form of `tf.Tensor`. However the gradients of keras Embedding layer are `tf.IndexedSlices`. 

[`tf.IndexedSlices`](https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/IndexedSlices) is used for sparse update.


Discussion:

Resolving patch:
diff --git a/elasticdl/python/common/ndarray.py b/elasticdl/python/common/ndarray.py
--- a/elasticdl/python/common/ndarray.py
+++ b/elasticdl/python/common/ndarray.py
@@ -1,4 +1,5 @@
 import numpy as np
+import tensorflow as tf
 
 from elasticdl.proto import elasticdl_pb2
 
@@ -23,15 +24,22 @@ def tensor_to_ndarray(tensor_pb):
             tensor_pb.dim,
             len(tensor_pb.content),
         )
-    arr = np.ndarray(
-        shape=tensor_pb.dim, dtype=np.float32, buffer=tensor_pb.content
-    )
+    if not tensor_pb.indices:
+        arr = np.ndarray(
+            shape=tensor_pb.dim, dtype=np.float32, buffer=tensor_pb.content
+        )
+    else:
+        values = np.ndarray(
+            shape=tensor_pb.dim, dtype=np.float32, buffer=tensor_pb.content
+        )
+        indices = tf.convert_to_tensor(tensor_pb.indices)
+        arr = tf.IndexedSlices(values, indices)
     tensor_pb.Clear()
 
     return arr
 
 
-def ndarray_to_tensor(arr):
+def ndarray_to_tensor(arr, indices=None):
     """Convert ndarray to Tensor PB"""
 
     if arr.dtype != np.float32:
@@ -41,5 +49,7 @@ def ndarray_to_tensor(arr):
     tensor = elasticdl_pb2.Tensor()
     tensor.dim.extend(arr.shape)
     tensor.content = arr.tobytes()
+    if indices:
+        tensor.indices.extend(indices)
 
     return tensor
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
@@ -36,6 +36,7 @@ def __init__(
         self._task_d = task_d
         self._lock = threading.Lock()
         self._gradient_sum = {}
+        self._gradient_sum_indexed = {}
         self._grad_to_wait = grads_to_wait
         self._grad_n = 0
         self._minibatch_size = minibatch_size
@@ -142,9 +143,14 @@ def _update_model(self):
         for k in self._gradient_sum:
             self._gradient_sum[k] = self._gradient_sum[k] / self._grad_to_wait
             grad_var.append((self._gradient_sum[k], self._model[k]))
+
+        for k in self._gradient_sum_indexed:
+            grad_var.append((self._gradient_sum_indexed[k], self._model[k]))
+
         self._opt.apply_gradients(grad_var)
         self._update_model_version()
         self._gradient_sum.clear()
+        self._gradient_sum_indexed.clear()
         self._grad_n = 0
 
     def get_model_version(self):
@@ -239,6 +245,7 @@ def ReportGradient(self, request, _):
         # TODO: Update task queue with task_id
         with self._lock:
             tmp = {}
+            indexed_grads = {}
             # Do sanity check before accumulating gradients.
             for k, v in request.gradient.items():
                 if k not in self._model:
@@ -246,11 +253,46 @@ def ReportGradient(self, request, _):
                         "Gradient key: %s is not part of model", k
                     )
                 arr = tensor_to_ndarray(v)
-                if arr.shape != self._model[k].numpy().shape:
-                    raise ValueError(
-                        "Gradient key: %s has incompatible dimension", k
+                if isinstance(arr, tf.IndexedSlices):
+                    if arr.values.shape[1] != self._model[k].numpy().shape[1]:
+                        raise ValueError(
+                            "Gradient key: %s has incompatible "
+                            "indexed slice dimension %d, expected %d"
+                            % (
+                                k,
+                                arr.values.shape[1],
+                                self._model[k].numpy().shape[1],
+                            )
+                        )
+
+                    max_index = tf.math.reduce_max(arr.indices).numpy()
+                    if max_index >= self._model[k].numpy().shape[0]:
+                        raise ValueError(
+                            "Gradient key: %s has wrong indices %d, "
+                            "out of range %d"
+                            % (
+                                k,
+                                max_index,
+                                self._model[k].numpy().shape[0] - 1,
+                            )
+                        )
+                    indexed_grads[k] = arr
+                else:
+                    if arr.shape != self._model[k].numpy().shape:
+                        raise ValueError(
+                            "Gradient key: %s has incompatible dimension", k
+                        )
+                    tmp[k] = arr
+
+            for k, v in indexed_grads.items():
+                if k not in self._gradient_sum_indexed:
+                    self._gradient_sum_indexed[k] = v
+                else:
+                    grads_s = self._gradient_sum_indexed[k]
+                    self._gradient_sum_indexed[k] = tf.IndexedSlices(
+                        tf.concat([grads_s.values, v.values], axis=0),
+                        tf.concat([grads_s.indices, v.indices], axis=0),
                     )
-                tmp[k] = arr
 
             for k, v in tmp.items():
                 if k in self._gradient_sum:
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
@@ -116,7 +116,14 @@ def report_gradient(self, grads):
         """
         req = elasticdl_pb2.ReportGradientRequest()
         for g, v in zip(grads, self._model.trainable_variables):
-            req.gradient[v.name].CopyFrom(ndarray_to_tensor(g.numpy()))
+            if isinstance(g, tf.IndexedSlices):
+                req.gradient[v.name].CopyFrom(
+                    ndarray_to_tensor(
+                        g.values.numpy(), tuple(g.indices.numpy())
+                    )
+                )
+            else:
+                req.gradient[v.name].CopyFrom(ndarray_to_tensor(g.numpy()))
         req.model_version = self._model_version
         res = self._stub.ReportGradient(req)
         return res.accepted, res.model_version


Test patch:
diff --git a/elasticdl/python/tests/indices_slices_gradient_test.py b/elasticdl/python/tests/indices_slices_gradient_test.py
new file mode 100644
--- /dev/null
+++ b/elasticdl/python/tests/indices_slices_gradient_test.py
@@ -0,0 +1,139 @@
+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 IndexedSlicesTest(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_indices_slices_correctness(self):
+        master, worker = self._create_master_and_worker()
+
+        values1 = tf.convert_to_tensor(
+            np.array([[1, 2, 3], [10, 11, 12]], dtype=np.float32)
+        )
+        values2 = tf.convert_to_tensor(
+            np.array([[7, 8, 9], [4, 5, 6]], dtype=np.float32)
+        )
+
+        indices1 = tf.convert_to_tensor([0, 3], dtype=tf.int64)
+        indices2 = tf.convert_to_tensor([2, 0], dtype=tf.int64)
+
+        grads1 = [
+            tf.IndexedSlices(values1, indices1),
+            tf.reshape(
+                tf.convert_to_tensor(range(1, 7), dtype=tf.float32),
+                shape=(6, 1),
+            ),
+        ]
+        grads2 = [
+            tf.IndexedSlices(values2, indices2),
+            tf.reshape(
+                tf.convert_to_tensor(range(13, 19), dtype=tf.float32),
+                shape=(6, 1),
+            ),
+        ]
+
+        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)
+            )
+        )
+
+        worker._model_version = 0
+        worker.report_gradient(grads1)
+        worker.report_gradient(grads2)
+
+        for i, j in zip(master._model.values(), expected_weights):
+            self.assertTrue(np.all(i.numpy() - j < 0.0001))
+
+    def test_wrong_indices(self):
+        master, worker = self._create_master_and_worker()
+
+        values = tf.zeros((2, 3))
+        indices = tf.convert_to_tensor([0, 4], dtype=tf.int64)
+        grads = [tf.IndexedSlices(values, indices), tf.zeros((6, 1))]
+        err_msg = ".*wrong indices %d, out of range %d" % (4, 3)
+        worker._model_version = 0
+        with self.assertRaisesRegex(ValueError, err_msg):
+            worker.report_gradient(grads)
+
+    def test_wrong_shape(self):
+        master, worker = self._create_master_and_worker()
+
+        values = tf.zeros((2, 4))
+        indices = tf.convert_to_tensor([0, 3], dtype=tf.int64)
+        grads = [tf.IndexedSlices(values, indices), tf.zeros((6, 1))]
+        err_msg = ".*incompatible indexed slice dimension %d, expected %d" % (
+            4,
+            3,
+        )
+        worker._model_version = 0
+        with self.assertRaisesRegex(ValueError, err_msg):
+            worker.report_gradient(grads)
+
+
+if __name__ == "__main__":
+    unittest.main()