Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fixing the inconsistency when exporting embeddings to cudf DataFrame #452

Merged
merged 5 commits into from
May 20, 2022
Merged

Fixing the inconsistency when exporting embeddings to cudf DataFrame #452

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
c41ae8f
Fixing the issue when exporting embeddings to cudf DataFrame
gabrielspmoreira May 18, 2022
e3d1f7a
Fixing test of embedding_table_dataset when GPU is not available
gabrielspmoreira May 18, 2022
249acf6
Fixing test
gabrielspmoreira May 18, 2022
1bed214
Merge branch 'main' into emb_export_fix
marcromeyn May 19, 2022
139dfc0
Merge branch 'main' into emb_export_fix
gabrielspmoreira May 19, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 9 additions & 2 deletions merlin/models/tf/features/embedding.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -311,8 +311,15 @@ def embedding_table_df(
embeddings = self.get_embedding_table(table_name, l2_normalization)
if gpu:
import cudf

df = cudf.from_dlpack(to_dlpack(tf.convert_to_tensor(embeddings)))
import cupy

# Note: It is not possible to convert Tensorflow tensors to the cudf dataframe
# directly using dlPack (as the example commented below) because cudf.from_dlpack()
# expects the 2D tensor to be in Fortran order (column-major), which is not
# supported by TF (https://github.com/rapidsai/cudf/issues/10754).
# df = cudf.from_dlpack(to_dlpack(tf.convert_to_tensor(embeddings)))
embeddings_cupy = cupy.fromDlpack(to_dlpack(tf.convert_to_tensor(embeddings)))
df = cudf.DataFrame(embeddings_cupy)
df.columns = [str(col) for col in list(df.columns)]
df.set_index(cudf.RangeIndex(0, embeddings.shape[0]))
else:
Expand Down
27 changes: 18 additions & 9 deletions tests/unit/tf/features/test_embedding.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -276,12 +276,6 @@ def test_embedding_features_yoochoose_pretrained_initializer(testing_data: Datas
assert np.allclose(emb_module.embedding_tables["categories"].numpy(), pretrained_emb_categories)


@pytest.mark.skip(
reason="This test fails because cudf.from_dlpack() we use internally in "
"EmbeddingFeatures().embedding_table_dataset() does not work properly. "
"This test should be enabled when the following cudf issue we reported is resolved: "
"https://github.com/rapidsai/cudf/issues/10754"
)
def test_embedding_features_exporting_and_loading_pretrained_initializer(testing_data: Dataset):
schema = testing_data.schema.select_by_tag(Tags.CATEGORICAL)
emb_module = mm.EmbeddingFeatures.from_schema(schema)
Expand All @@ -290,15 +284,30 @@ def test_embedding_features_exporting_and_loading_pretrained_initializer(testing
_ = emb_module(mm.sample_batch(testing_data, batch_size=10, include_targets=False))
item_id_embeddings = emb_module.embedding_tables["item_id"]

items_embeddings_dataset = emb_module.embedding_table_dataset(Tags.ITEM_ID)
items_embeddings_dataset = emb_module.embedding_table_dataset(Tags.ITEM_ID, gpu=False)
assert np.allclose(
item_id_embeddings.numpy(), items_embeddings_dataset.to_ddf().compute().to_pandas().values
item_id_embeddings.numpy(), items_embeddings_dataset.to_ddf().compute().values
)

emb_init = mm.TensorInitializer.from_dataset(items_embeddings_dataset)

assert np.allclose(item_id_embeddings.numpy(), emb_init(item_id_embeddings.shape).numpy())

# Test GPU export if available
try:
import cudf # noqa: F401

items_embeddings_dataset = emb_module.embedding_table_dataset(Tags.ITEM_ID, gpu=True)
assert np.allclose(
item_id_embeddings.numpy(),
items_embeddings_dataset.to_ddf().compute().to_pandas().values,
)

emb_init = mm.TensorInitializer.from_dataset(items_embeddings_dataset)
assert np.allclose(item_id_embeddings.numpy(), emb_init(item_id_embeddings.shape).numpy())

except ImportError:
pass


def test_shared_embeddings(music_streaming_data: Dataset):
inputs = mm.InputBlock(music_streaming_data.schema)
Expand Down