Repository: sql-machine-learning/elasticdl

Issue:
[master]Use tf.ResourceVariable to store model
Currently we store model as a <string, ndarray> map. when using tf.optimizer.apply_gradient() to update model, we need to convert the map to ResourceVariable and back. It is better to change model to a <string, ResourceVariable> map to avoid copy and conversion.


Discussion:

Resolving patch:
diff --git a/elasticdl/master/servicer.py b/elasticdl/master/servicer.py
--- a/elasticdl/master/servicer.py
+++ b/elasticdl/master/servicer.py
@@ -1,5 +1,7 @@
 import threading
+import numpy as np
 
+import tensorflow as tf
 from proto import master_pb2
 from proto import master_pb2_grpc
 from util.converter import NdarrayToTensor, TensorToNdarray
@@ -12,12 +14,21 @@ def __init__(self, logger, grads_to_wait):
         self.logger = logger
         self._lock = threading.Lock()
         # TODO: random initialization
+        # A <string, tf.ResourceVariable> map. We use tf.ResourceVariable
+        # instead ndarray to avoid copying and conversion when calling
+        # optimizer's apply_gradients() function.
         self._model = {}
         self._version = 0
         self._gradient_sum = {}
         self._grad_to_wait = grads_to_wait
         self._grad_n = 0
 
+    def _set_model_var(self, name, value):
+        """Add or set model variable. Value should be a float32 ndarray"""
+        if value.dtype != np.float32:
+            raise ValueError("Value should be a float32 numpy array")
+        self._model[name] = tf.Variable(value, name=name, use_resource=True)
+
     def GetTask(self, request, context):
         # TODO: implent task queues. Return an empty task for now.
         res = master_pb2.Task()
@@ -38,7 +49,7 @@ def GetModel(self, request, context):
         with self._lock:
             res.version = self._version
             for k, v in self._model.items():
-                res.param[k].CopyFrom(NdarrayToTensor(v))
+                res.param[k].CopyFrom(NdarrayToTensor(v.numpy()))
         return res
 
     def ReportTaskResult(self, request, context):
@@ -76,7 +87,7 @@ def ReportTaskResult(self, request, context):
                         "Gradient key: %s is not part of model", k
                     )
                 arr = TensorToNdarray(v)
-                if arr.shape != self._model[k].shape:
+                if arr.shape != self._model[k].numpy().shape:
                     raise ValueError(
                         "Gradient key: %s has incompatible dimension", k
                     )


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
@@ -2,6 +2,9 @@
 import unittest
 import numpy as np
 
+import tensorflow as tf
+tf.enable_eager_execution()
+
 from google.protobuf import empty_pb2
 
 from proto import master_pb2
@@ -29,7 +32,7 @@ def testGetModel(self):
         req.min_version = 0
 
         # Get version 0
-        master._model["x"] = np.array([1.0, 1.0], dtype=np.float32)
+        master._set_model_var("x", np.array([1.0, 1.0], dtype=np.float32))
         model = master.GetModel(req, None)
         self.assertEqual(0, model.version)
         self.assertEqual(["x"], list(model.param.keys()))
@@ -39,8 +42,8 @@ def testGetModel(self):
 
         # increase master's model version, now should get version 1
         master._version = 1
-        master._model["x"] = np.array([2.0, 2.0], dtype=np.float32)
-        master._model["y"] = np.array([12.0, 13.0], dtype=np.float32)
+        master._set_model_var("x", np.array([2.0, 2.0], dtype=np.float32))
+        master._set_model_var("y", np.array([12.0, 13.0], dtype=np.float32))
         model = master.GetModel(req, None)
         self.assertEqual(1, model.version)
         self.assertEqual(["x", "y"], list(model.param.keys()))
@@ -70,8 +73,8 @@ def makeGrad():
 
         master = MasterServicer(logging.getLogger(), 3)
         master._version = 1
-        master._model["x"] = np.array([2.0], dtype=np.float32)
-        master._model["y"] = np.array([12.0, 13.0], dtype=np.float32)
+        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))
 
         # Report a future version, should raise exception
         req = makeGrad()