add multi layer DCN as a building block for DCN V2 (DLRM_DCN) model a…

…nd add the option to not concat dense before top stack PiperOrigin-RevId: 601465868
tensorflow · Feb 16, 2024 · 5f0a39f · 5f0a39f
1 parent 1e45dd4
commit 5f0a39f
Show file tree

Hide file tree

Showing 49 changed files with 351 additions and 58 deletions.
diff --git a/setup.py b/setup.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/__init__.py b/tensorflow_recommenders/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/examples/__init__.py b/tensorflow_recommenders/examples/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/examples/movielens.py b/tensorflow_recommenders/examples/movielens.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/examples/nbtool.py b/tensorflow_recommenders/examples/nbtool.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/__init__.py b/tensorflow_recommenders/experimental/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/layers/__init__.py b/tensorflow_recommenders/experimental/layers/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/layers/embedding/__init__.py b/tensorflow_recommenders/experimental/layers/embedding/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/layers/embedding/partial_tpu_embedding.py b/tensorflow_recommenders/experimental/layers/embedding/partial_tpu_embedding.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/layers/embedding/partial_tpu_embedding_test.py b/tensorflow_recommenders/experimental/layers/embedding/partial_tpu_embedding_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/models/__init__.py b/tensorflow_recommenders/experimental/models/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/models/ranking.py b/tensorflow_recommenders/experimental/models/ranking.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -61,6 +61,7 @@ def __init__(
       bottom_stack: Optional[tf.keras.layers.Layer] = None,
       feature_interaction: Optional[tf.keras.layers.Layer] = None,
       top_stack: Optional[tf.keras.layers.Layer] = None,
+      concat_dense: bool = True,
       task: Optional[tasks.Task] = None) -> None:
     """Initializes the model.
 
@@ -80,20 +81,32 @@ def __init__(
       top_stack: The `top_stack` layer is applied to the `feature_interaction`
         output. The output of top_stack should be in the range [0, 1]. If it is
         None, MLP with layer sizes [512, 256, 1] is used.
+      concat_dense: Weather to concatenate the interaction output with dense
+        embedding vector again before feeding into the top stack
       task: The task which the model should optimize for. Defaults to a
-        `tfrs.tasks.Ranking` task with a binary cross-entropy loss, suitable
-        for tasks like click prediction.
+        `tfrs.tasks.Ranking` task with a binary cross-entropy loss, suitable for
+        tasks like click prediction.
     """
 
     super().__init__()
 
     self._embedding_layer = embedding_layer
-    self._bottom_stack = bottom_stack if bottom_stack else layers.blocks.MLP(
-        units=[256, 64, 16], final_activation="relu")
-    self._top_stack = top_stack if top_stack else layers.blocks.MLP(
-        units=[512, 256, 1], final_activation="sigmoid")
-    self._feature_interaction = (feature_interaction if feature_interaction
-                                 else feature_interaction_lib.DotInteraction())
+    self._concat_dense = concat_dense
+    self._bottom_stack = (
+        bottom_stack
+        if bottom_stack
+        else layers.blocks.MLP(units=[256, 64, 16], final_activation="relu")
+    )
+    self._top_stack = (
+        top_stack
+        if top_stack
+        else layers.blocks.MLP(units=[512, 256, 1], final_activation="sigmoid")
+    )
+    self._feature_interaction = (
+        feature_interaction
+        if feature_interaction
+        else feature_interaction_lib.DotInteraction()
+    )
 
     if task is not None:
       self._task = task
@@ -211,8 +224,12 @@ def call(self, inputs: Dict[str, tf.Tensor]) -> tf.Tensor:
 
     interaction_args = sparse_embedding_vecs + [dense_embedding_vec]
     interaction_output = self._feature_interaction(interaction_args)
-    feature_interaction_output = tf.concat(
-        [dense_embedding_vec, interaction_output], axis=1)
+    if self._concat_dense:
+      feature_interaction_output = tf.concat(
+          [dense_embedding_vec, interaction_output], axis=1
+      )
+    else:
+      feature_interaction_output = interaction_output
 
     prediction = self._top_stack(feature_interaction_output)
 

diff --git a/tensorflow_recommenders/experimental/models/ranking_test.py b/tensorflow_recommenders/experimental/models/ranking_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -127,6 +127,8 @@ class RankingTest(tf.test.TestCase, parameterized.TestCase):
           # Top stack.
           (lambda: None, lambda: tfrs.layers.blocks.MLP(
               units=[40, 20, 1], final_activation="sigmoid")),
+          # Concat Dense.
+          (True, False),
           # Use weights.
           (True, False),
           # Size threshold.
@@ -135,6 +137,7 @@ def test_ranking_model(self,
                          feature_interaction_layer,
                          bottom_stack,
                          top_stack,
+                         concat_dense=True,
                          use_weights=False,
                          size_threshold=10):
     """Tests a ranking model."""
@@ -151,7 +154,8 @@ def test_ranking_model(self,
             size_threshold=size_threshold),
         bottom_stack=bottom_stack(),
         feature_interaction=feature_interaction_layer(),
-        top_stack=top_stack())
+        top_stack=top_stack(),
+        concat_dense=concat_dense)
     model.compile(optimizer=optimizer, steps_per_execution=5)
 
     dataset = _generate_synthetic_data(

diff --git a/tensorflow_recommenders/experimental/optimizers/__init__.py b/tensorflow_recommenders/experimental/optimizers/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/optimizers/clippy_adagrad.py b/tensorflow_recommenders/experimental/optimizers/clippy_adagrad.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/optimizers/clippy_adagrad_test.py b/tensorflow_recommenders/experimental/optimizers/clippy_adagrad_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/optimizers/composite_optimizer.py b/tensorflow_recommenders/experimental/optimizers/composite_optimizer.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/optimizers/composite_optimizer_test.py b/tensorflow_recommenders/experimental/optimizers/composite_optimizer_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/__init__.py b/tensorflow_recommenders/layers/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/blocks.py b/tensorflow_recommenders/layers/blocks.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/embedding/__init__.py b/tensorflow_recommenders/layers/embedding/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/embedding/tpu_embedding_layer.py b/tensorflow_recommenders/layers/embedding/tpu_embedding_layer.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/embedding/tpu_embedding_layer_test.py b/tensorflow_recommenders/layers/embedding/tpu_embedding_layer_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/factorized_top_k.py b/tensorflow_recommenders/layers/factorized_top_k.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/factorized_top_k_test.py b/tensorflow_recommenders/layers/factorized_top_k_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/feature_interaction/__init__.py b/tensorflow_recommenders/layers/feature_interaction/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -16,3 +16,4 @@
 
 from tensorflow_recommenders.layers.feature_interaction.dcn import Cross
 from tensorflow_recommenders.layers.feature_interaction.dot_interaction import DotInteraction
+from tensorflow_recommenders.layers.feature_interaction.multi_layer_dcn import MultiLayerDCN
diff --git a/tensorflow_recommenders/layers/feature_interaction/dcn.py b/tensorflow_recommenders/layers/feature_interaction/dcn.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/feature_interaction/dcn_test.py b/tensorflow_recommenders/layers/feature_interaction/dcn_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/feature_interaction/dot_interaction.py b/tensorflow_recommenders/layers/feature_interaction/dot_interaction.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/feature_interaction/dot_interaction_test.py b/tensorflow_recommenders/layers/feature_interaction/dot_interaction_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Recommenders Authors.
+# Copyright 2024 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.