Repository: sql-machine-learning/elasticdl

Issue:
[master] PS model update 



Discussion:
For now, we use tf.training.optimizer's apply_gradients and rely on optimizer's algorithm.

Resolving patch:
diff --git a/elasticdl/master/servicer.py b/elasticdl/master/servicer.py
--- a/elasticdl/master/servicer.py
+++ b/elasticdl/master/servicer.py
@@ -2,6 +2,8 @@
 import numpy as np
 
 import tensorflow as tf
+assert tf.executing_eagerly()
+
 from proto import master_pb2
 from proto import master_pb2_grpc
 from util.ndarray import ndarray_to_tensor, tensor_to_ndarray
@@ -10,8 +12,9 @@
 class MasterServicer(master_pb2_grpc.MasterServicer):
     """Master service implementation"""
 
-    def __init__(self, logger, grads_to_wait):
+    def __init__(self, logger, grads_to_wait, optimizer):
         self.logger = logger
+        self._opt = optimizer
         self._lock = threading.Lock()
         # TODO: random initialization
         # A <string, tf.ResourceVariable> map. We use tf.ResourceVariable
@@ -52,6 +55,17 @@ def GetModel(self, request, context):
                 res.param[k].CopyFrom(ndarray_to_tensor(v.numpy()))
         return res
 
+    def _update_model(self):
+        assert self._lock.locked()
+        grad_var = []
+        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]))
+        self._opt.apply_gradients(grad_var)
+        self._version += 1
+        self._gradient_sum.clear()
+        self._grad_n = 0
+
     def ReportTaskResult(self, request, context):
         if request.model_version > self._version:
             err_msg = "Model version %d out of range, current version: %d" % (
@@ -101,10 +115,7 @@ def ReportTaskResult(self, request, context):
 
             self._grad_n += 1
             if self._grad_n >= self._grad_to_wait:
-                # TODO: update model
-                self._version += 1
-                self._gradient_sum.clear()
-                self._grad_n = 0
+                self._update_model()
         res.accepted = True
         res.model_version = self._version
         return res


Test patch:
diff --git a/elasticdl/master/servicer_test.py b/elasticdl/master/servicer_test.py
--- a/elasticdl/master/servicer_test.py
+++ b/elasticdl/master/servicer_test.py
@@ -3,6 +3,7 @@
 import numpy as np
 
 import tensorflow as tf
+
 tf.enable_eager_execution()
 
 from google.protobuf import empty_pb2
@@ -14,7 +15,7 @@
 
 class ServicerTest(unittest.TestCase):
     def testGetTask(self):
-        master = MasterServicer(logging.getLogger(), 2)
+        master = MasterServicer(logging.getLogger(), 2, None)
 
         # No task yet, make sure the returned versions are as expected.
         task = master.GetTask(empty_pb2.Empty(), None)
@@ -27,7 +28,7 @@ def testGetTask(self):
         self.assertEqual(1, task.model_version)
 
     def testGetModel(self):
-        master = MasterServicer(logging.getLogger(), 2)
+        master = MasterServicer(logging.getLogger(), 2, None)
         req = master_pb2.GetModelRequest()
         req.min_version = 0
 
@@ -66,12 +67,14 @@ def makeGrad():
                 ndarray_to_tensor(np.array([0.1], dtype=np.float32))
             )
             req.gradient["y"].CopyFrom(
-                ndarray_to_tensor(np.array([0.1, 0.2], dtype=np.float32))
+                ndarray_to_tensor(np.array([0.03, 0.06], dtype=np.float32))
             )
             req.model_version = 1
             return req
 
-        master = MasterServicer(logging.getLogger(), 3)
+        master = MasterServicer(
+            logging.getLogger(), 3, tf.train.GradientDescentOptimizer(0.1)
+        )
         master._version = 1
         master._set_model_var("x", np.array([2.0], dtype=np.float32))
         master._set_model_var("y", np.array([12.0, 13.0], dtype=np.float32))
@@ -126,7 +129,7 @@ def makeGrad():
             np.array([0.2], dtype=np.float32), master._gradient_sum["x"]
         )
         np.testing.assert_array_equal(
-            np.array([0.1, 0.2], dtype=np.float32), master._gradient_sum["y"]
+            np.array([0.03, 0.06], dtype=np.float32), master._gradient_sum["y"]
         )
         self.assertEqual(2, master._grad_n)
 
@@ -136,6 +139,15 @@ def makeGrad():
         res = master.ReportTaskResult(req, None)
         self.assertTrue(res.accepted)
         self.assertEqual(2, res.model_version)
-        # TODO: verify model when model updating is in place.
         self.assertFalse(master._gradient_sum)
         self.assertEqual(0, master._grad_n)
+        np.testing.assert_array_equal(
+            # [2] - 0.1 * [0.1]
+            np.array([1.99], dtype=np.float32),
+            master._model["x"].numpy(),
+        )
+        np.testing.assert_array_equal(
+            # [12, 13] - 0.1 * [0.02, 0.04]
+            np.array([11.998, 12.996], dtype=np.float32),
+            master._model["y"].numpy(),
+        )