From 9b241a4abd1c96fd6290ea0f9765b7303688b2e1 Mon Sep 17 00:00:00 2001 From: Julian Qian Date: Wed, 2 Oct 2024 13:54:43 -0700 Subject: [PATCH 1/2] [fix] ensure shape in the safe_embedding_lookup --- .../dynamic_embedding_ops_test.py | 31 +++++++++++++++++ .../kernel_tests/redis_table_ops_test.py | 5 ++- .../kernel_tests/redis_table_variable_test.py | 6 +++- .../ops/distributed_embedding_variable.py | 3 +- .../python/ops/dynamic_embedding_ops.py | 34 ++++++++++++------- .../python/ops/dynamic_embedding_variable.py | 5 +-- .../python/ops/embedding_weights.py | 4 ++- .../python/ops/ragged_embedding_ops.py | 2 +- .../python/ops/shadow_embedding_ops.py | 16 ++++++--- 9 files changed, 82 insertions(+), 24 deletions(-) diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_ops_test.py b/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_ops_test.py index cb60c3063..8723daa04 100644 --- a/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_ops_test.py +++ b/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_ops_test.py @@ -71,8 +71,10 @@ try: from tf_keras.initializers import Zeros from tf_keras.optimizers import Adam + from tf_keras import Input, Model, layers except: from tensorflow.keras.initializers import Zeros + from tensorflow.keras import Input, Model, layers try: from tensorflow.keras.optimizers import Adam except: @@ -567,6 +569,35 @@ def test_scope_reuse_sparse_embedding_lookup(self): self.assertAllEqual(p1_reuse._tables[0].name, "test_p1_mht_1of1") self.assertAllEqual(p2._tables[0].name, "test_p2_mht_1of1") + def test_keras_input_safe_sparse_embedding_lookup(self): + embeddings = de.get_variable(name="p1") + + @tf.function + def model(indices_input, ids_input): + shape_input = tf.constant([2, 3, 4], dtype=tf.int64) + + sparse_ids = tf.sparse.SparseTensor(indices=indices_input, + values=ids_input, + dense_shape=shape_input) + + embeddings_result = de.safe_embedding_lookup_sparse( + embeddings, sparse_ids, name="safe_sp_emb", return_trainable=False) + return embeddings_result + + indices_data = tf.constant([ + [0, 0, 0], + [0, 0, 1], + [0, 0, 2], + [0, 1, 0], + [1, 0, 0], + [1, 1, 0], + [1, 1, 1], + ], + dtype=tf.int64) + ids_data = tf.constant([0, 1, -1, -1, 2, 0, 1], dtype=tf.int64) + + # output = model(indices_data, ids_data) + def test_scope_reuse_safe_sparse_embedding_lookup(self): indices = [ [0, 0, 0], diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/redis_table_ops_test.py b/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/redis_table_ops_test.py index 585dc8568..d6c7c2ec0 100755 --- a/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/redis_table_ops_test.py +++ b/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/redis_table_ops_test.py @@ -38,7 +38,10 @@ from tensorflow.python.framework import ops from tensorflow.python.framework import sparse_tensor from tensorflow.python.framework import test_util -from tensorflow.keras import layers +try: + from tensorflow.keras import layers +except: + from tf_keras import layers from tensorflow.python.keras import optimizer_v2 from tensorflow.python.ops import array_ops from tensorflow.python.ops import init_ops diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/redis_table_variable_test.py b/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/redis_table_variable_test.py index 951483de9..baf8af53d 100644 --- a/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/redis_table_variable_test.py +++ b/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/redis_table_variable_test.py @@ -34,7 +34,11 @@ from tensorflow_recommenders_addons.utils.check_platform import is_windows, is_macos, is_arm64, is_linux, is_raspi_arm from tensorflow.core.protobuf import config_pb2 -from tensorflow.keras import layers +try: + from tensorflow.keras import layers +except: + from tf_keras import layers + from tensorflow.python.client import session from tensorflow.python.eager import context from tensorflow.python.framework import constant_op diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/ops/distributed_embedding_variable.py b/tensorflow_recommenders_addons/dynamic_embedding/python/ops/distributed_embedding_variable.py index d07c36d34..353cc3088 100644 --- a/tensorflow_recommenders_addons/dynamic_embedding/python/ops/distributed_embedding_variable.py +++ b/tensorflow_recommenders_addons/dynamic_embedding/python/ops/distributed_embedding_variable.py @@ -19,6 +19,7 @@ def verify_embedding_weights(self, sparse_ids, sparse_weights=None): def embedding_lookup(self, ids, name=None, - max_norm=None) -> (tf.Tensor, EmbeddingWeights): + max_norm=None, + return_trainable=False) -> (tf.Tensor, EmbeddingWeights): raise NotImplementedError("embedding_lookup is not supported in " "DistributedVariableWrapper") \ No newline at end of file diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_ops.py b/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_ops.py index 1b1755854..5c55e4297 100644 --- a/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_ops.py +++ b/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_ops.py @@ -94,14 +94,17 @@ def embedding_lookup_unique(params, ids_flat = array_ops.reshape(ids, math_ops.reduce_prod(shape, keepdims=True)) unique_ids, idx = array_ops.unique(ids_flat) - unique_embeddings, trainable_ = de.embedding_lookup( - params, - unique_ids, - partition_strategy=partition_strategy, - name=name, - validate_indices=None, - max_norm=validate_indices, - return_trainable=True) + result = de.embedding_lookup(params, + unique_ids, + partition_strategy=partition_strategy, + name=name, + validate_indices=None, + max_norm=validate_indices, + return_trainable=return_trainable) + if return_trainable: + unique_embeddings, trainable_ = result + else: + unique_embeddings = result embeddings_flat = array_ops.gather(unique_embeddings, idx) embeddings_shape = array_ops.concat( [shape, array_ops.shape(unique_embeddings)[1:]], 0) @@ -216,9 +219,12 @@ def embedding_lookup_sparse( ids = sp_ids.values ids, idx = array_ops.unique(ids) - embeddings, trainable_ = params.embedding_lookup(ids, - name=name, - max_norm=max_norm) + embeddings = params.embedding_lookup(ids, + name=name, + max_norm=max_norm, + return_trainable=return_trainable) + if return_trainable: + embeddings, trainable_ = embeddings if embeddings.dtype in (dtypes.float16, dtypes.bfloat16): embeddings = math_ops.cast(embeddings, dtypes.float32) @@ -372,7 +378,7 @@ def safe_embedding_lookup_sparse( if sparse_weights is not None: sparse_weights, _ = de.math.sparse_fill_empty_rows(sparse_weights, 1.0) - result, trainable_ = embedding_lookup_sparse( + result = embedding_lookup_sparse( embedding_weights, sparse_ids, sparse_weights, @@ -380,8 +386,10 @@ def safe_embedding_lookup_sparse( partition_strategy=partition_strategy, name=name + "/embedding_lookup_sparse", max_norm=max_norm, - return_trainable=True, + return_trainable=return_trainable, ) + if (return_trainable): + result, trainable_ = result if default_id is None: # Broadcast is_row_empty to the same shape as embedding_lookup_result, diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_variable.py b/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_variable.py index 7e411c9a7..9e9d08ac5 100644 --- a/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_variable.py +++ b/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_variable.py @@ -692,13 +692,14 @@ def verify_embedding_weights(self, sparse_ids, sparse_weights=None): def embedding_lookup(self, ids, name=None, - max_norm=None) -> (tf.Tensor, EmbeddingWeights): + max_norm=None, + return_trainable=False) -> (tf.Tensor, EmbeddingWeights): return embedding_lookup( self, ids, name=name + '/embedding_lookup', max_norm=max_norm, - return_trainable=True, + return_trainable=return_trainable, ) @property diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/ops/embedding_weights.py b/tensorflow_recommenders_addons/dynamic_embedding/python/ops/embedding_weights.py index 39402dd09..2d600f28e 100644 --- a/tensorflow_recommenders_addons/dynamic_embedding/python/ops/embedding_weights.py +++ b/tensorflow_recommenders_addons/dynamic_embedding/python/ops/embedding_weights.py @@ -57,7 +57,9 @@ def verify_embedding_weights(self, sparse_ids, sparse_weights=None): def embedding_lookup(self, ids, name=None, - max_norm=None) -> (tf.Tensor, "EmbeddingWeights"): + max_norm=None, + return_trainable=False + ) -> (tf.Tensor, "EmbeddingWeights"): """ embedding lookup, and store the result. No by-product will be introduced in this call. So it can be decorated by `tf.function`. diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/ops/ragged_embedding_ops.py b/tensorflow_recommenders_addons/dynamic_embedding/python/ops/ragged_embedding_ops.py index bf1c8c8ef..26213acf3 100644 --- a/tensorflow_recommenders_addons/dynamic_embedding/python/ops/ragged_embedding_ops.py +++ b/tensorflow_recommenders_addons/dynamic_embedding/python/ops/ragged_embedding_ops.py @@ -140,7 +140,7 @@ def _embedding_lookup_sparse_impl( segment_ids = ops.convert_to_tensor(segment_ids, name="segment_ids") ids, idx = array_ops.unique(ids) - embeddings, _ = params.embedding_lookup(ids, name=name) + embeddings = params.embedding_lookup(ids, name=name) if not ignore_weights: if segment_ids.dtype != dtypes.int32: segment_ids = math_ops.cast(segment_ids, dtypes.int32) diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/ops/shadow_embedding_ops.py b/tensorflow_recommenders_addons/dynamic_embedding/python/ops/shadow_embedding_ops.py index 0f1f44526..72b496e0a 100644 --- a/tensorflow_recommenders_addons/dynamic_embedding/python/ops/shadow_embedding_ops.py +++ b/tensorflow_recommenders_addons/dynamic_embedding/python/ops/shadow_embedding_ops.py @@ -168,8 +168,12 @@ def verify_embedding_weights(self, sparse_ids, sparse_weights=None): def embedding_lookup(self, ids, name=None, - max_norm=None) -> (tf.Tensor, EmbeddingWeights): - return embedding_lookup(self, ids, name), self + max_norm=None, + return_trainable=False) -> (tf.Tensor, EmbeddingWeights): + if return_trainable: + return embedding_lookup(self, ids, name, False), self + else: + return embedding_lookup(self, ids, name) def prefetch_values(self, update=False): if self.params.bp_v2: @@ -443,5 +447,9 @@ def __alltoall_embedding_lookup__(self, ids): def embedding_lookup(self, ids, name=None, - max_norm=None) -> (tf.Tensor, EmbeddingWeights): - return self.__alltoall_embedding_lookup__(ids), self.shadow + max_norm=None, + return_trainable=False) -> (tf.Tensor, EmbeddingWeights): + if return_trainable: + return self.__alltoall_embedding_lookup__(ids), self.shadow + else: + return self.__alltoall_embedding_lookup__(ids) From b08954d4e21ff41ce07840f689b4b7e7acf92e6b Mon Sep 17 00:00:00 2001 From: Julian Qian Date: Wed, 2 Oct 2024 18:08:56 -0700 Subject: [PATCH 2/2] [fix] clean disk space to avoid out of disk problem in the ci --- .github/workflows/make_wheel_Linux_x86.sh | 2 + .../dynamic_embedding_ops_test.py | 65 ++++++++++++++----- tools/testing/build_and_run_tests.sh | 9 ++- 3 files changed, 57 insertions(+), 19 deletions(-) diff --git a/.github/workflows/make_wheel_Linux_x86.sh b/.github/workflows/make_wheel_Linux_x86.sh index 5a9a5ce90..69de8a043 100644 --- a/.github/workflows/make_wheel_Linux_x86.sh +++ b/.github/workflows/make_wheel_Linux_x86.sh @@ -56,6 +56,8 @@ if [[ "$TF_VERSION" =~ ^2\.1[3-9]\.[0-9]$ ]] ; then export PROTOBUF_VERSION='4.23.4' fi +docker system prune -f + DOCKER_BUILDKIT=1 docker build --no-cache \ -f tools/docker/build_wheel.Dockerfile \ --output type=local,dest=wheelhouse \ diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_ops_test.py b/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_ops_test.py index 8723daa04..5f4792523 100644 --- a/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_ops_test.py +++ b/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_ops_test.py @@ -81,6 +81,52 @@ from tensorflow.keras.legacy.optimizers import Adam +class MyModel(layers.Layer): + + def __init__(self): + super().__init__() + self.embeddings = de.get_variable(name="p1") + + @tf.function + def call(self, indices_input, ids_input): + shape_input = tf.constant([2, 3, 4], dtype=tf.int64) + + sparse_ids = tf.sparse.SparseTensor(indices=indices_input, + values=ids_input, + dense_shape=shape_input) + + embeddings_result = de.safe_embedding_lookup_sparse(self.embeddings, + sparse_ids, + name="safe_sp_emb", + return_trainable=False) + return embeddings_result + + +# class MyModel(tf.Module): +# def __init__(self): +# super().__init__() +# # Create the variable as an attribute of the module +# self.embeddings = de.get_variable(name="p1") +# +# @tf.function +# def __call__(self, indices_input, ids_input): +# shape_input = tf.constant([2, 3, 4], dtype=tf.int64) +# +# sparse_ids = tf.sparse.SparseTensor( +# indices=indices_input, +# values=ids_input, +# dense_shape=shape_input +# ) +# +# embeddings_result = de.safe_embedding_lookup_sparse( +# self.embeddings, +# sparse_ids, +# name="safe_sp_emb", +# return_trainable=False +# ) +# return embeddings_result + + # pylint: disable=missing-class-docstring # pylint: disable=missing-function-docstring def _type_converter(tf_type): @@ -570,20 +616,6 @@ def test_scope_reuse_sparse_embedding_lookup(self): self.assertAllEqual(p2._tables[0].name, "test_p2_mht_1of1") def test_keras_input_safe_sparse_embedding_lookup(self): - embeddings = de.get_variable(name="p1") - - @tf.function - def model(indices_input, ids_input): - shape_input = tf.constant([2, 3, 4], dtype=tf.int64) - - sparse_ids = tf.sparse.SparseTensor(indices=indices_input, - values=ids_input, - dense_shape=shape_input) - - embeddings_result = de.safe_embedding_lookup_sparse( - embeddings, sparse_ids, name="safe_sp_emb", return_trainable=False) - return embeddings_result - indices_data = tf.constant([ [0, 0, 0], [0, 0, 1], @@ -593,9 +625,10 @@ def model(indices_input, ids_input): [1, 1, 0], [1, 1, 1], ], + name="indices_data", dtype=tf.int64) - ids_data = tf.constant([0, 1, -1, -1, 2, 0, 1], dtype=tf.int64) - + ids_data = tf.constant([0, 1, -1, -1, 2, 0, 1], name="ids", dtype=tf.int64) + model = MyModel() # output = model(indices_data, ids_data) def test_scope_reuse_safe_sparse_embedding_lookup(self): diff --git a/tools/testing/build_and_run_tests.sh b/tools/testing/build_and_run_tests.sh index 59cdb74ca..9d2a6553f 100644 --- a/tools/testing/build_and_run_tests.sh +++ b/tools/testing/build_and_run_tests.sh @@ -31,6 +31,8 @@ if [ "$TF_VERSION" = "2.6.3" ] ; then python -m pip install numpy==1.19.5 --force-reinstall fi +bazel clean --expunge + # Avoid SystemError: initialization of _pywrap_checkpoint_reader raised unreported exception pip install tensorflow==$TF_VERSION @@ -79,8 +81,9 @@ fi python -m pytest -v -s --functions-durations=20 --modules-durations=5 $IGNORE_HKV $SKIP_CUSTOM_OP_TESTS_FLAG $EXTRA_ARGS ./tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/ - - # Release disk space bazel clean --expunge -rm -f ./tensorflow_recommenders_addons/dynamic_embedding/core/_*_ops.so +sudo rm -f ./tensorflow_recommenders_addons/dynamic_embedding/core/_*_ops.so +sudo rm -rf /tmp/* +apt-get clean && rm -rf /var/lib/apt/lists/* +sudo rm -rf /var/log/* \ No newline at end of file