udpate

Author: jatinsharechat

README.md

@@ -165,10 +165,8 @@ cd recommenders-addons
 
 # This script links project with TensorFlow dependency
 python configure.py
-
 bazel build --enable_runfiles build_pip_pkg
 bazel-bin/build_pip_pkg artifacts
-
 pip install artifacts/tensorflow_recommenders_addons-*.whl
 ```
 #### GPU Support

demo/dynamic_embedding/movielens-1m-keras-with-horovod/movielens-1m-keras-with-horovod.py

@@ -23,7 +23,8 @@
 
 
 def has_horovod() -> bool:
-  return 'OMPI_COMM_WORLD_RANK' in os.environ or 'PMI_RANK' in os.environ
+  #return 'OMPI_COMM_WORLD_RANK' in os.environ or 'PMI_RANK' in os.environ
+  return True
 
 
 def config():
@@ -275,6 +276,7 @@ def __init__(self,
         init_capacity=init_capacity,
         kv_creator=kv_creator_dense,
         short_file_name=True,
+        restrict_policy=de.TimestampRestrictPolicy,
     )
 
     kv_creator_sparse = get_kv_creator(mpi_size, mpi_rank, init_capacity,
@@ -290,6 +292,7 @@ def __init__(self,
         init_capacity=init_capacity,
         kv_creator=kv_creator_sparse,
         short_file_name=True,
+        restrict_policy=de.TimestampRestrictPolicy,
     )
 
     self.dnn = tf.keras.layers.Dense(
@@ -648,7 +651,7 @@ def train():
     # horovod callback is used to broadcast the value generated by initializer of rank0.
     hvd_opt_init_callback = de.keras.callbacks.DEHvdBroadcastGlobalVariablesCallback(
         root_rank=0)
-    callbacks_list = [hvd_opt_init_callback, ckpt_callback]
+    callbacks_list = [hvd_opt_init_callback]#, ckpt_callback]
   else:
     callbacks_list = [ckpt_callback]
 
@@ -664,6 +667,8 @@ def train():
             steps_per_epoch=FLAGS.steps_per_epoch,
             verbose=1 if get_rank() == 0 else 0)
 
+  print(model.user_embedding.sparse_embedding_layer.params.restrict_policy)
+
   export_to_savedmodel(model, FLAGS.model_dir)
   export_for_serving(model, FLAGS.export_dir)
 

tensorflow_recommenders_addons/dynamic_embedding/python/keras/models.py

@@ -121,11 +121,11 @@ def _save_de_var(de_var, proc_size=1, proc_rank=0):
                                  proc_size=proc_size,
                                  proc_rank=proc_rank)
 
-  def _maybe_save_restrict_policy_params(var, proc_size=1, proc_rank=0):
-    if not hasattr(var, "restrict_policy"):
+  def _maybe_save_restrict_policy_params(de_var, proc_size=1, proc_rank=0):
+    if not hasattr(de_var, "restrict_policy"):
       return
-    if var.restrict_policy is not None:
-      de_var = var.restrict_policy._restrict_var
+    if de_var.restrict_policy is not None:
+      de_var = de_var.restrict_policy._restrict_var
       _save_de_var(de_var, proc_size=proc_size, proc_rank=proc_rank)
 
   def _traverse_emb_layers_and_save(proc_size=1, proc_rank=0):
@@ -136,7 +136,7 @@ def _traverse_emb_layers_and_save(proc_size=1, proc_rank=0):
         continue
       de_var = var.params
       _save_de_var(de_var, proc_size=proc_size, proc_rank=proc_rank)
-      _maybe_save_restrict_policy_params(var, proc_size=proc_size, proc_rank=proc_rank)
+      _maybe_save_restrict_policy_params(de_var, proc_size=proc_size, proc_rank=proc_rank)
 
   if hvd is None:
     call_original_save_func()