update

Author: jatinsharechat

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

@@ -29,12 +29,13 @@ def has_horovod() -> bool:
 def config():
   # callback calls hvd.rank() so we need to initialize horovod here
   hvd.init()
+  print("Size: ", hvd.size())
   if has_horovod():
     print("Horovod is enabled.")
     if hvd.rank() > 0:
       os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
     # Horovod: pin GPU to be used to process local rank (one GPU per process)
-    config_gpu(hvd.local_rank())
+    # config_gpu(hvd.local_rank())
   else:
     config_gpu()
 
@@ -460,8 +461,8 @@ def call(self, features):
 def get_dataset(batch_size=1):
   ds = tfds.load("movielens/1m-ratings",
                  split="train",
-                 data_dir="~/dataset",
-                 download=True)
+                 # data_dir="~/dataset",
+                 download=False)
   features = ds.map(
       lambda x: {
           "movie_id":

demo/dynamic_embedding/movielens-1m-keras-with-horovod/start.sh

@@ -1,6 +1,7 @@
 #!/bin/bash
 rm -rf ./export_dir
-gpu_num=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
+#gpu_num=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
+gpu_num=5
 export gpu_num
 horovodrun -np $gpu_num python movielens-1m-keras-with-horovod.py --mode="train" --model_dir="./model_dir" --export_dir="./export_dir" \
-  --steps_per_epoch=${1:-20000} --shuffle=${2:-True}
+  --steps_per_epoch=${1:-10} --shuffle=${2:-False}

tensorflow_recommenders_addons/dynamic_embedding/python/keras/models.py

@@ -122,6 +122,8 @@ def _save_de_var(de_var, proc_size=1, proc_rank=0):
                                  proc_rank=proc_rank)
 
   def _maybe_save_restrict_policy_params(var, proc_size=1, proc_rank=0):
+    if not hasattr(var, "restrict_policy"):
+      return
     if var.restrict_policy is not None:
       de_var = var.restrict_policy._restrict_var
       _save_de_var(de_var, proc_size=proc_size, proc_rank=proc_rank)