resource var

Author: jq

demo/dynamic_embedding/movielens-1m-keras-ps/movielens-1m-keras-ps.py

@@ -72,13 +72,13 @@ def __init__(self,
         user_embedding_size,
         initializer=embedding_initializer,
         devices=self.devices,
-        with_unique=False,
+        # with_unique=False,
         name='user_embedding')
     self.movie_embedding = de.keras.layers.SquashedEmbedding(
         movie_embedding_size,
         initializer=embedding_initializer,
         devices=self.devices,
-        with_unique=False,
+        # with_unique=False,
         name='movie_embedding')
 
     self.dnn1 = tf.keras.layers.Dense(
@@ -105,10 +105,11 @@ def __init__(self,
   @tf.function
   def call(self, features):
     user_id = tf.reshape(features['user_id'], (-1, 1))
-    movie_id = tf.reshape(features['movie_id'], (-1, 1))
+    # movie_id = tf.reshape(features['movie_id'], (-1, 1))
     user_latent = self.user_embedding(user_id)
-    movie_latent = self.movie_embedding(movie_id)
-    latent = tf.concat([user_latent, movie_latent], axis=1)
+    # movie_latent = self.movie_embedding(movie_id)
+    # latent = tf.concat([user_latent, movie_latent], axis=1)
+    latent = user_latent
     x = self.dnn1(latent)
     x = self.dnn2(x)
     x = self.dnn3(x)
@@ -129,7 +130,7 @@ def __init__(self, strategy, train_bs, test_bs, epochs, steps_per_epoch,
         "/job:ps/replica:0/task:{}/device:CPU:0".format(idx)
         for idx in range(self.num_ps)
     ]
-    self.embedding_size = 32
+    self.embedding_size = 4
     self.train_bs = train_bs
     self.test_bs = test_bs
     self.epochs = epochs
@@ -254,10 +255,10 @@ def start_chief(config):
       cluster_spec, task_type="chief", task_id=0)
   strategy = tf_dist.experimental.ParameterServerStrategy(cluster_resolver)
   runner = Runner(strategy=strategy,
-                  train_bs=64,
+                  train_bs=4,
                   test_bs=1,
                   epochs=1,
-                  steps_per_epoch=1000,
+                  steps_per_epoch=4,
                   model_dir=None,
                   export_dir=None)
   runner.train()

tensorflow_recommenders_addons/dynamic_embedding/python/keras/layers/embedding.py

@@ -303,6 +303,9 @@ def call(self, ids):
     Returns:
       A embedding output with shape (shape(ids), embedding_size).
     """
+    tfprint = tf.print("ids_8a:", ids, output_stream=tf.compat.v1.logging.error)
+    with tf.control_dependencies([tfprint]):
+      pass
     return de.shadow_ops.embedding_lookup_unique(self.shadow, ids,
                                                  self.embedding_size,
                                                  self.with_unique, self.name)

tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_optimizer.py

@@ -16,6 +16,8 @@
 """patch on optimizers"""
 
 import functools
+
+import tensorflow as tf
 from packaging import version
 import six
 
@@ -184,6 +186,9 @@ def apply_grad_to_update_var(var, grad):
                   "Cannot use a constraint function on a sparse variable.")
             if "apply_state" in self._sparse_apply_args:
               apply_kwargs["apply_state"] = apply_state
+            # printop = tf.print("ids_8d:", output_stream=tf.compat.v1.logging.error)
+            # with tf.control_dependencies([printop]):
+            #   pass
             with ops.control_dependencies(_before):
               _apply_op = self._resource_apply_sparse_duplicate_indices(
                   grad.values, var, grad.indices, **apply_kwargs)

tensorflow_recommenders_addons/dynamic_embedding/python/ops/parameter_server.py

@@ -0,0 +1,37 @@
+from tensorflow.python.distribute import ps_values, distribute_lib
+from tensorflow.python.distribute.distribute_lib import _get_per_thread_mode
+from tensorflow.python.distribute.parameter_server_strategy_v2 import ParameterServerStrategyV2, \
+  ParameterServerStrategyV2Extended
+from tensorflow.python.ops import variables
+import tensorflow as tf
+
+
+class DEPerWorkerVariable(ps_values.PerWorkerVariable):
+  def __init__(self, *args, **kwargs):
+    super(DEPerWorkerVariable, self).__init__(*args, **kwargs)
+
+def create_per_worker_de_variable(strategy, name, dtype, shape):
+  # printop = tf.print("st_2:", strategy,
+  #                    tf.distribute.get_replica_context() ,
+  #                    output_stream=tf.compat.v1.logging.error)
+  # with tf.control_dependencies([printop]):
+    with strategy.scope():
+      return variables.Variable(initial_value=(),
+        shape=shape, dtype=dtype, name=name,
+        per_worker_de_variable=True)
+
+original_create_variable = ParameterServerStrategyV2Extended._create_variable
+
+def patched_create_variable(self, next_creator, **kwargs):
+  if kwargs.pop("per_worker_de_variable", False):
+    return _create_per_worker_de_variable(self, next_creator, **kwargs)
+  return original_create_variable(self, next_creator, **kwargs)
+
+def _create_per_worker_de_variable(strategy_extended, next_creator, **kwargs):
+  return DEPerWorkerVariable(strategy_extended._container_strategy(), next_creator, **kwargs)
+
+ParameterServerStrategyV2Extended._create_variable = patched_create_variable
+
+class DEParameterServerStrategy(ParameterServerStrategyV2):
+  def __init__(self, cluster_resolver, variable_partitioner=None):
+    super(DEParameterServerStrategy, self).__init__(cluster_resolver, variable_partitioner)

tensorflow_recommenders_addons/dynamic_embedding/python/ops/shadow_embedding_ops.py

@@ -37,7 +37,7 @@
 
 import tensorflow as tf
 
-from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import distribute_lib, ps_values
 from tensorflow.python.eager import context
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
@@ -49,6 +49,10 @@
 from tensorflow_recommenders_addons import dynamic_embedding as de
 from tensorflow_recommenders_addons.dynamic_embedding.python.ops.embedding_weights import EmbeddingWeights, \
   TrainableWrapper
+from tensorflow_recommenders_addons.dynamic_embedding.python.ops.parameter_server import create_per_worker_de_variable, \
+  DEPerWorkerVariable
+from tensorflow_recommenders_addons.dynamic_embedding.python.train.utils import is_parameter_server_strategy
+from tensorflow.python.ops import variables
 
 if version.parse(tf.__version__) >= version.parse("2.10"):
   from tensorflow.python.trackable import base as trackable
@@ -114,15 +118,18 @@ def __init__(self,
     ids_name = self._name + '-ids'
     if ids is None:
       self.ids = get_de_resource_variable(
-          trainable=False,
           collections=collections,
           name=ids_name,
           dtype=self.params.key_dtype,
           distribute_strategy=distribute_strategy,
           shape=tensor_shape.TensorShape(None))
     else:
       if not isinstance(ids, resource_variable_ops.ResourceVariable):
-        raise TypeError('If ids is set, it needs to be a ResourceVariable')
+        tfprint = tf.print("ids_8c:", ids, type(ids), ids.__class__.__name__, output_stream=tf.compat.v1.logging.error)
+        with tf.control_dependencies([tfprint]):
+          pass
+      #     not isinstance(ids, variables.Variable)):
+      #   raise TypeError('If ids is set, it needs to be a ResourceVariable or ps_values.PerWorkerVariable')
       self.ids = ids
 
     model_mode = kwargs.get('model_mode', None)
@@ -152,7 +159,6 @@ def __init__(self,
     exists_name = self._name + '-exists'
     if exists is None:
       self.exists = get_de_resource_variable(
-          trainable=False,
           collections=collections,
           name=exists_name,
           dtype=dtypes.bool,
@@ -272,10 +278,14 @@ def embedding_lookup(
   with ops.name_scope(name, "shadow_embedding_lookup"):
     with ops.colocate_with(None, ignore_existing=True):
       if de.ModelMode.CURRENT_SETTING == de.ModelMode.TRAIN:
+        tfprint = tf.print("ids_8b:", shadow_.ids, ids, output_stream=tf.compat.v1.logging.error)
+        with tf.control_dependencies([tfprint]):
+          pass
         with ops.control_dependencies([shadow_._reset_ids(ids)]):
           result = shadow_.read_value(do_prefetch=True)
       else:
         result = shadow_.params.lookup(ids)
+
       return result
 
 
@@ -360,14 +370,17 @@ def __init__(self, *args, **kwargs):
     super(DEResourceVariable, self).__init__(*args, **kwargs)
 
 
-def get_de_resource_variable(trainable,
+def get_de_resource_variable(
                              collections,
                              name,
                              dtype,
                              distribute_strategy,
                              shape=tensor_shape.TensorShape(None)):
-  return DEResourceVariable((),
-                            trainable=trainable,
+  if is_parameter_server_strategy(distribute_strategy):
+    return create_per_worker_de_variable(distribute_strategy, name, dtype, shape)
+  else:
+    return DEResourceVariable((),
+                            trainable=False,
                             collections=collections,
                             name=name,
                             dtype=dtype,
@@ -377,7 +390,7 @@ def get_de_resource_variable(trainable,
 
 def is_de_resource_variable(var):
   return isinstance(var, DEResourceVariable) or isinstance(
-      var, TrainableWrapper)
+      var, TrainableWrapper) or isinstance(var, DEPerWorkerVariable)
 
 
 class HvdVariable(EmbeddingWeights):