Repository: sql-machine-learning/elasticdl

Issue:
Optimizer for Embedding Layer Step 2: Set Values to Variables
According to [design doc](https://github.com/wangkuiyi/elasticdl/blob/develop/elasticdl/doc/distributed_embedding_layer_design.md#optimizer), set slot values to optimizers and transform embedding values form `tf.Tensor` to `tf.Variable`.



Discussion:

Resolving patch:
diff --git a/elasticdl/python/master/optimizer_wrapper.py b/elasticdl/python/master/optimizer_wrapper.py
--- a/elasticdl/python/master/optimizer_wrapper.py
+++ b/elasticdl/python/master/optimizer_wrapper.py
@@ -61,6 +61,32 @@ def _get_slot_name_from_key(key):
     return key.split("-")[-2]
 
 
+# This function is taken from `tensorflow.keras.optimizers.Optimizer._var_key`.
+# https://github.com/tensorflow/tensorflow/blob/71d73e56a2e66e4a6805d967cfa48ea
+# 594f8c54e/tensorflow/python/keras/optimizer_v2/optimizer_v2.py#L1033
+def _var_key(var):
+    """Key for representing a primary variable, for looking up slots.
+
+    In graph mode the name is derived from the var shared name.
+    In eager mode the name is derived from the var unique id.
+    If distribution strategy exists, get the primary variable first.
+
+    Arguments:
+        var: the variable.
+
+    Returns:
+        the unique name of the variable.
+    """
+
+    # pylint: disable=protected-access
+    # Get the distributed variable if it exists.
+    if getattr(var, "_distributed_container", None) is not None:
+        var = var._distributed_container()
+    if var._in_graph_mode:
+        return var._shared_name
+    return var._unique_id
+
+
 class OptimizerWrapper(object):
     """ ElasticDL optimizer wrapper.
 
@@ -83,8 +109,10 @@ def __init__(self, opt, kv_store_endpoint, embedding_dims):
         """
         self._opt = opt
         self._kv_store_endpoint = kv_store_endpoint
-        self._embedding_dims = embedding_dims
+        self._embed_dims = embedding_dims
         self._slot_initial_value = {}
+        self._embed_variables = {}
+        self._slot_variables = {}
 
         # TODO: support more TensorFlow optimizers
         # "-" in slot name is not supported
@@ -123,7 +151,7 @@ def apply_gradients(self, grads_and_vars):
         grads_and_vars_local = []
         grads_and_vars_kv_store = []
         for grad, var in grads_and_vars:
-            layer_name = _get_embedding_layer_name_from_var(grad, var)
+            layer_name = _get_embedding_layer_name_from_var(var)
             if layer_name:
                 grads_and_vars_kv_store.append((grad, layer_name))
             else:
@@ -135,10 +163,14 @@ def apply_gradients(self, grads_and_vars):
             grads_and_vars_kv_store
         )
 
+        self._set_embedding_values_to_variables(
+            grads_and_vars_kv_store, embed_values
+        )
+        self._set_slot_values_to_variables(slot_values)
+
         # TODO: implement the following logic:
-        # 1. set embedding values and slot values to TensorFlow Variables
-        # 2. call self._opt.apply_gradients
-        # 3. report updated values to Redis
+        # * call self._opt.apply_gradients
+        # * report updated values to Redis
 
     def _lookup_embeddings_and_slots(self, grads_and_vars):
         """Look up embedding vectors and slot values form kv store.
@@ -256,10 +288,125 @@ def _generate_lookup_keys(self, grads_and_vars):
 
     def _initialize_unknown_slot(self, layer_name, slot_name):
         """Initialize unknown slot."""
-        slot_dim = self._embedding_dims[layer_name]
+        slot_dim = self._embed_dims[layer_name]
         initial_value = self._slot_initial_value[slot_name]
         return np.full((slot_dim,), initial_value, np.float32)
 
+    def _set_embedding_values_to_variables(self, grads_and_vars, values):
+        """Set embedding values to embedding variables."""
+        for i, (grad, layer_name) in enumerate(grads_and_vars):
+            value = values[layer_name]
+            variable = self._get_embedding_variable(layer_name)
+            if variable is None:
+                variable = self._create_embedding_variable(layer_name, value)
+            else:
+                variable.assign(value)
+            grads_and_vars[i] = (grad, variable)
+
+    def _set_slot_values_to_variables(self, values):
+        """Set slot values to slot variables in TensorFlow optimizers."""
+        for layer_name, slots in values.items():
+            for slot_name, slot_value in slots.items():
+                # `variable` points to the variable object saved in
+                # TensorFlow optimizer, i.e. self._opt
+                variable = self._get_slot_variable(layer_name, slot_name)
+                if variable is None:
+                    self._create_slot_variable(
+                        layer_name, slot_name, slot_value
+                    )
+                else:
+                    variable.assign(slot_value)
+
+    def _get_slot_variable(self, layer_name, slot_name):
+        """Get the variable for specified slot."""
+        return self._slot_variables.get(layer_name, {}).get(slot_name, None)
+
+    def _get_embedding_variable(self, layer_name):
+        """Get the variable for the specified ElasticDL embedding layer."""
+        return self._embed_variables.get(layer_name, None)
+
+    def _create_embedding_variable(self, layer_name, initial_value):
+        """Create a variable for an ElasticDL embedding layer."""
+        dim = self._embed_dims[layer_name]
+        # Use shape `(None, dim)` for embedding variable because `shape[0]`
+        # equals to the number of unqiue ids in the minibatch data, and
+        # this number may differ between different iterations
+        shape = tf.TensorShape((None, dim))
+
+        if self._embed_variables.get(layer_name, None) is not None:
+            raise RuntimeError(
+                "Embedding variable with layer name=%s has already be "
+                "created." % (layer_name)
+            )
+
+        embed_var = tf.Variable(
+            initial_value,
+            name=layer_name,
+            shape=shape,
+            dtype=tf.float32,
+            trainable=False,
+        )
+        self._embed_variables[layer_name] = embed_var
+        return embed_var
+
+    def _create_slot_variable(self, layer_name, slot_name, initial_value):
+        """Create a variable for the specified slot."""
+        dim = self._embed_dims[layer_name]
+        # Use shape `(None, dim)` for slot variable because `shape[0]`
+        # equals to the number of unqiue ids in the minibatch data, and
+        # this number may differ between different iterations
+        shape = tf.TensorShape((None, dim))
+
+        slot_variables_dict = self._slot_variables.setdefault(layer_name, {})
+        if slot_variables_dict.get(slot_name, None) is not None:
+            raise RuntimeError(
+                "Slot variable with (layer name=%s, slot name=%s) has "
+                "already be created." % (layer_name, slot_name)
+            )
+
+        embed_var = self._get_embedding_variable(layer_name)
+        if embed_var is None:
+            raise RuntimeError(
+                "Embedding variable for layer %s should be already created."
+                % (layer_name)
+            )
+        slot_var = self._create_slot_variable_in_optimizer(
+            embed_var, slot_name, shape, initial_value
+        )
+        slot_variables_dict[slot_name] = slot_var
+        return slot_var
+
+    # This is a function modified from TensorFlow optimizers.
+    # https://github.com/tensorflow/tensorflow/blob/
+    # 69b1feac62276edcc509ac88af229c6236e645fe/tensorflow/python
+    # /keras/optimizer_v2/optimizer_v2.py#L567
+    def _create_slot_variable_in_optimizer(
+        self, embed_var, slot_name, shape, initial_value
+    ):
+        """Create variable for a slot and save it in TensorFlow optimizer."""
+        if slot_name not in self._opt._slot_names:
+            self._opt._slot_names.append(slot_name)
+        var_key = _var_key(embed_var)
+        slot_dict = self._opt._slots.setdefault(var_key, {})
+        slot_var = slot_dict.get(slot_name, None)
+        if slot_var is None:
+            slot_var_name = "%s/%s" % (embed_var._shared_name, slot_name)
+            slot_var = self._opt.add_weight(
+                name=slot_var_name,
+                shape=shape,
+                dtype=embed_var.dtype,
+                initializer=initial_value,
+                trainable=False,
+            )
+            slot_dict[slot_name] = slot_var
+            self._opt._weights.append(slot_var)
+            return slot_var
+        else:
+            raise RuntimeError(
+                "Variable with var_key %s and slot_name %s is not expected to "
+                "be in self._opt." % (var_key, slot_name)
+            )
+
     @property
     def allowed_slot_names(self):
         return self._allowed_slot_names


Test patch:
diff --git a/elasticdl/python/tests/optimizer_wrapper_test.py b/elasticdl/python/tests/optimizer_wrapper_test.py
--- a/elasticdl/python/tests/optimizer_wrapper_test.py
+++ b/elasticdl/python/tests/optimizer_wrapper_test.py
@@ -178,6 +178,111 @@ def test_lookup(self):
                     (slot_values[layer][slot] - values < 0.0001).all()
                 )
 
+    def test_set_slot_values_to_variables(self):
+        layers = ["test-1", "test-2"]
+        slots = ["m", "v"]
+        id_num = 3
+        embedding_dims = {layer: 4 for layer in layers}
+        all_values = np.arange(48).reshape(12, 4).astype(np.float32)
+
+        slot_values = {}
+        offset = 0
+        for layer in layers:
+            for slot in slots:
+                start = offset
+                end = offset + id_num
+                slot_values.setdefault(layer, {}).setdefault(
+                    slot, all_values[start:end]
+                )
+                offset = end
+
+        opt = Adam()
+        opt_wrapper = OptimizerWrapper(opt, None, embedding_dims)
+        for layer in layers:
+            opt_wrapper._create_embedding_variable(layer, tf.zeros((1, 4)))
+        opt_wrapper._set_slot_values_to_variables(slot_values)
+        self.assertTrue(len(opt.weights) == 4)
+        for layer in layers:
+            slots_dict = None
+            for k, v in opt._slots.items():
+                if k.startswith(layer):
+                    slots_dict = v
+                    break
+
+            for slot in slots:
+                self.assertTrue(
+                    (
+                        slots_dict[slot].numpy() == slot_values[layer][slot]
+                    ).all()
+                )
+                self.assertTrue(
+                    slots_dict[slot]
+                    == opt_wrapper._slot_variables[layer][slot]
+                )
+
+                slots_dict[slot].assign(tf.ones((10, 4)))
+                self.assertTrue(
+                    (
+                        opt_wrapper._slot_variables[layer][slot].numpy() - 1.0
+                        < 0.0001
+                    ).all()
+                )
+                opt_wrapper._slot_variables[layer][slot].assign(
+                    -tf.ones((10, 4))
+                )
+                self.assertTrue(
+                    (slots_dict[slot].numpy() + 1.0 < 0.0001).all()
+                )
+
+        slot_values_new = {"test-1": {"m": np.zeros((3, 4), np.float32)}}
+        opt_wrapper._set_slot_values_to_variables(slot_values_new)
+        self.assertTrue(
+            (opt_wrapper._slot_variables["test-1"]["m"].numpy() < 0.0001).all()
+        )
+
+    def test_set_embedding_values_to_variables(self):
+        layers = ["test-1", "test-2"]
+        id_num = 3
+        embedding_dims = {layer: 4 for layer in layers}
+        all_values = np.arange(16).reshape(4, 4)
+
+        embedding_values = {}
+        offset = 0
+        for layer in layers:
+            start = offset
+            end = offset + id_num
+            embedding_values.setdefault(layer, all_values[start:end])
+            offset = end
+
+        opt = SGD()
+        opt_wrapper = OptimizerWrapper(opt, None, embedding_dims)
+        grads_and_vars = [
+            ("test-1-grads", "test-1"),
+            ("test-2-grads", "test-2"),
+        ]
+        opt_wrapper._set_embedding_values_to_variables(
+            grads_and_vars, embedding_values
+        )
+        for i, layer in enumerate(layers):
+            self.assertTrue(
+                (
+                    opt_wrapper._embed_variables[layer].numpy()
+                    == embedding_values[layer]
+                ).all()
+            )
+            self.assertTrue(
+                grads_and_vars[i][1] == opt_wrapper._embed_variables[layer]
+            )
+
+        embedding_values_new = {"test-1": np.zeros((3, 4), np.float32)}
+        grads_and_vars = [("test-1-grads", "test-1")]
+        opt_wrapper._set_embedding_values_to_variables(
+            grads_and_vars, embedding_values_new
+        )
+        self.assertTrue(
+            (opt_wrapper._embed_variables["test-1"].numpy() < 0.0001).all()
+        )
+
 
 if __name__ == "__main__":
     unittest.main()