Repository: sql-machine-learning/elasticdl

Issue:
[PS-1] Create `Parameters` and `EmbeddingTable` class definition
According to [PS design](https://github.com/sql-machine-learning/elasticdl/blob/develop/docs/designs/ps_design.md#data-structure).

ps.main also create a `Parameters` instance to store model parameters

Two steps:

- [x] add `EmbeddingTable` #1305 
- [x] add `Parameters` #1359


Discussion:

Resolving patch:
diff --git a/elasticdl/python/ps/parameters.py b/elasticdl/python/ps/parameters.py
new file mode 100644
--- /dev/null
+++ b/elasticdl/python/ps/parameters.py
@@ -0,0 +1,56 @@
+import tensorflow as tf
+
+from elasticdl.python.common.ndarray import tensor_to_ndarray
+from elasticdl.python.ps.embedding_table import create_embedding_table
+
+
+class Parameters(object):
+    """
+    There are two kinds of parameters:
+
+    1. non-embedding parameters, or dense tensor parameters. We save it
+       in a hashmap `non-embedding_params`, the key is the parameter name,
+       the value is a tf.Variable` object.
+    2. embedding parameters, or row-sparse parameters. We save it in a
+       hashmap `embedding_params`, the key is the embedding layer name,
+       the value is an `EmbeddingTable` object.
+
+    """
+
+    def __init__(self):
+        self.init_status = False
+        self.non_embedding_params = {}
+        self.embedding_params = {}
+
+    def get_embedding_param(self, name, indices):
+        if name not in self.embedding_params:
+            raise ValueError(
+                "Please initialize embedding param %s first!", name
+            )
+        return self.embedding_params[name].get(indices)
+
+    def set_embedding_param(self, name, indices, values):
+        if name not in self.embedding_params:
+            raise ValueError(
+                "Please initialize embedding param %s first!", name
+            )
+        self.embedding_params[name].set(indices, values)
+
+    def init_from_model_pb(self, model_pb):
+        if not self.init_status:
+            tensors_pb = model_pb.param
+            embeddings_pb = model_pb.embedding_table_info
+            self._init_non_embedding_params(tensors_pb)
+            self._init_embedding_params(embeddings_pb)
+            self.init_status = True
+
+    def _init_non_embedding_params(self, tensors_pb):
+        for pb in tensors_pb:
+            name = pb.name
+            arr = tensor_to_ndarray(pb)
+            var = tf.Variable(name=name, initial_value=arr, trainable=True)
+            self.non_embedding_params[name] = var
+
+    def _init_embedding_params(self, embeddings_pb):
+        for pb in embeddings_pb:
+            self.embedding_params[pb.name] = create_embedding_table(pb)


Test patch:
diff --git a/elasticdl/python/tests/parameters_test.py b/elasticdl/python/tests/parameters_test.py
new file mode 100644
--- /dev/null
+++ b/elasticdl/python/tests/parameters_test.py
@@ -0,0 +1,108 @@
+import unittest
+
+import numpy as np
+import tensorflow as tf
+
+from elasticdl.proto.elasticdl_pb2 import EmbeddingTableInfo, Model
+from elasticdl.python.common.ndarray import ndarray_to_tensor
+from elasticdl.python.ps.parameters import Parameters
+
+
+class ParametersTest(unittest.TestCase):
+    def setUp(self):
+        self.params = Parameters()
+
+        self.model_pb = Model()
+        self.tensors_pb = self.model_pb.param
+        self.embeddings_pb = self.model_pb.embedding_table_info
+
+        arr1 = np.random.uniform(size=(3, 4))
+        tensor1_pb = ndarray_to_tensor(arr1, "x")
+        arr2 = np.random.uniform(size=(4, 5))
+        tensor2_pb = ndarray_to_tensor(arr2, "y")
+        self.tensors_pb.extend([tensor1_pb, tensor2_pb])
+
+        self.embedding_table_name = "embedding_1"
+        self.embedding_dim = 10
+        embedding_pb = EmbeddingTableInfo()
+        embedding_pb.name = self.embedding_table_name
+        embedding_pb.dim = self.embedding_dim
+        embedding_pb.initializer = "uniform"
+
+        self.embeddings_pb.append(embedding_pb)
+
+    def _clear_params(self):
+        self.params.non_embedding_params.clear()
+        self.params.embedding_params.clear()
+
+    def _test_get_embedding_param(self):
+        indices = [0, 3, 7]
+
+        res = self.params.get_embedding_param(
+            self.embedding_table_name, indices
+        )
+        self.assertTupleEqual(res.shape, (3, 10))
+
+        res = self.params.get_embedding_param(self.embedding_table_name, [])
+        self.assertIsNone(res)
+
+        with self.assertRaises(ValueError):
+            self.params.get_embedding_param("tom", indices)
+
+    def test_init_from_model_pb(self):
+        self._clear_params()
+        self.params.init_from_model_pb(self.model_pb)
+
+        res = self.params.non_embedding_params
+        self.assertTrue("x" in res)
+        self.assertTrue("y" in res)
+        self.assertTrue(res["x"].trainable)
+        self.assertTupleEqual(tuple(res["y"].shape.as_list()), (4, 5))
+
+        self._test_get_embedding_param()
+
+    def test_non_embedding_params(self):
+        self._clear_params()
+
+        res = self.params.non_embedding_params
+        self.assertFalse(any(res))
+
+        variables = {
+            "x": tf.Variable(1, name="x"),
+            "y": tf.Variable(2, name="y"),
+        }
+
+        self.params.non_embedding_params = variables
+        self.assertTrue("x" in self.params.non_embedding_params)
+        self.assertTrue("y" in self.params.non_embedding_params)
+
+    def test_get_embedding_param(self):
+        self._clear_params()
+        self.params._init_embedding_params(self.embeddings_pb)
+        self._test_get_embedding_param()
+
+    def test_set_embedding_param(self):
+        self._clear_params()
+        self.params._init_embedding_params(self.embeddings_pb)
+        indices = [100, 34, 8]
+        x = len(indices)
+        values = np.random.uniform(size=x * self.embedding_dim).reshape(
+            (x, self.embedding_dim)
+        )
+
+        self.params.set_embedding_param(
+            self.embedding_table_name, indices, values
+        )
+
+        row0 = self.params.get_embedding_param(
+            self.embedding_table_name, [100]
+        )
+        row1 = self.params.get_embedding_param(self.embedding_table_name, [34])
+        row2 = self.params.get_embedding_param(self.embedding_table_name, [8])
+
+        rows = [row0, row1, row2]
+        rows = np.concatenate(rows)
+        np.testing.assert_array_equal(rows, values)
+
+        with self.assertRaises(ValueError):
+            self.params.set_embedding_param("tom", [0, 1, 2], values)