fix gpu_num

tensorflow · May 27, 2024 · 431b8b8 · 431b8b8
1 parent 0b9a5ee
commit 431b8b8
Show file tree

Hide file tree

Showing 7 changed files with 18 additions and 18 deletions.
diff --git a/.github/workflows/make_wheel_macOS_x86.sh b/.github/workflows/make_wheel_macOS_x86.sh
@@ -40,7 +40,4 @@ bazel build \
 
 bazel-bin/build_pip_pkg artifacts $NIGHTLY_FLAG
 
-# Setting DYLD_LIBRARY_PATH to help delocate finding tensorflow after the rpath invalidation
-export DYLD_LIBRARY_PATH=$DYLD_LIBRARY_PATH:$(python -c 'import configure; print(configure.get_tf_shared_lib_dir())')
 delocate-wheel -w wheelhouse -v --ignore-missing-dependencies artifacts/*.whl
-
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
@@ -111,7 +111,7 @@ jobs:
     strategy:
       matrix:
         # TODO: add back 'windows-latest' when it can be compiled.
-        os: ['self-hosted-gpu']
+        os: ['ubuntu-20.04']
         py-version: ['3.7', '3.8', '3.9', '3.10', '3.11']
         tf-version: ['2.11.0', '2.15.1']
         tf-need-cuda: ['1']

diff --git a/demo/dynamic_embedding/movielens-1m-keras-with-horovod/README.md b/demo/dynamic_embedding/movielens-1m-keras-with-horovod/README.md
@@ -3,7 +3,7 @@
 - dataset: [movielen/1m-ratings](https://www.tensorflow.org/datasets/catalog/movielens#movielens1m-ratings)
 - model: DNN
 - Running mode: Graph mode and Keras by using Horovod AllToAll Embedding as model parameters parallelism
-
+- enable gpu by `python3 -m pip install tensorflow[and-cuda]`
 ## start train:
 By default, this shell will start a train task with N workers as GPU number on local machine.
 

diff --git a/demo/dynamic_embedding/movielens-1m-keras-with-horovod/movielens-1m-keras-with-horovod.py b/demo/dynamic_embedding/movielens-1m-keras-with-horovod/movielens-1m-keras-with-horovod.py
@@ -63,11 +63,13 @@
         ],
                       dtype=tf.int32,
                       name='user_occupation_label'),
-    'raw_user_age':
+    'bucketized_user_age':
         tf.TensorSpec(shape=[
             None,
             1,
-        ], dtype=tf.int32, name='raw_user_age'),
+        ],
+                      dtype=tf.int32,
+                      name='bucketized_user_age'),
     'movie_id':
         tf.TensorSpec(shape=[
             None,
@@ -126,7 +128,7 @@
         'input_tensor': None,
         'pretreated_tensor': None,
     },
-    'raw_user_age': {
+    'bucketized_user_age': {
         'code': 106,
         'dtype': tf.int32,
         'dim': 1,
@@ -398,8 +400,8 @@ def call(self, features):
 def get_dataset(batch_size=1):
   ds = tfds.load("movielens/1m-ratings",
                  split="train",
-                 data_dir="/dataset",
-                 download=False)
+                 data_dir="~/dataset",
+                 download=True)
   features = ds.map(
       lambda x: {
           "movie_id":
@@ -412,8 +414,8 @@ def get_dataset(batch_size=1):
               tf.cast(x["user_gender"], tf.int32),
           "user_occupation_label":
               tf.cast(x["user_occupation_label"], tf.int32),
-          "raw_user_age":
-              tf.cast(x["raw_user_age"], tf.int32),
+          "bucketized_user_age":
+              tf.cast(x["bucketized_user_age"], tf.int32),
           "timestamp":
               tf.cast(x["timestamp"] - 880000000, tf.int32),
       })

diff --git a/demo/dynamic_embedding/movielens-1m-keras-with-horovod/start.sh b/demo/dynamic_embedding/movielens-1m-keras-with-horovod/start.sh
@@ -1,4 +1,5 @@
 #!/bin/bash
 rm -rf ./export_dir
-gpu_num = nvidia-smi --query-gpu=name --format=csv,noheader | wc -l
-horovodrun -np $gpu_num python movielens-1m-keras-with-horovod.py --mode="train" --model_dir="./model_dir" --export_dir="./export_dir" 
+gpu_num=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
+export gpu_num
+horovodrun -np $gpu_num python movielens-1m-keras-with-horovod.py --mode="train" --model_dir="./model_dir" --export_dir="./export_dir"
diff --git a/pytest.ini b/pytest.ini
diff --git a/tools/testing/build_and_run_tests.sh b/tools/testing/build_and_run_tests.sh
@@ -56,7 +56,10 @@ if [ "$(uname)" != "Darwin" ]; then
   bash /install/install_horovod.sh $HOROVOD_VERSION --only-cpu
 fi
 # TODO(jamesrong): Test on GPU.
-CUDA_VISIBLE_DEVICES="" mpirun -np 2 -H localhost:2 --allow-run-as-root pytest -v ./tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/horovod_sync_train_test.py
+CUDA_VISIBLE_DEVICES=""
+if [ "$TF_NEED_CUDA" -eq 1 ]; then
+  mpirun -np 2 -H localhost:2 --allow-run-as-root pytest -v ./tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/horovod_sync_train_test.py
+fi
 # Reinstall Horovod after tests
 if [ "$(uname)" != "Darwin" ]; then
   # Mac only with MPI