tfrs.experimental.models.Ranking

A configurable ranking model.

Inherits From: Model

tfrs.experimental.models.Ranking(
    embedding_layer: tf.keras.layers.Layer,
    bottom_stack: Optional[tf.keras.layers.Layer] = None,
    feature_interaction: Optional[tf.keras.layers.Layer] = None,
    top_stack: Optional[tf.keras.layers.Layer] = None,
    task: Optional[tfrs.tasks.Task] = None
) -> None

This class represents a sensible and reasonably flexible configuration for a ranking model that can be used for tasks such as CTR prediction.

It can be customized as needed, and its constituent blocks can be changed by passing user-defined alternatives.

For example:

  • Pass feature_interaction = tfrs.layers.feature_interaction.DotInteraction() to train a DLRM model, or pass

    feature_interaction = tf.keras.Sequential([
  tf.keras.layers.Concatenate(),
  tfrs.layers.feature_interaction.Cross()
])

    to train a DCN model.

  • Pass task = tfrs.tasks.Ranking(loss=tf.keras.losses.BinaryCrossentropy()) to train a CTR prediction model, and tfrs.tasks.Ranking(loss=tf.keras.losses.MeanSquaredError()) to train a rating prediction model.

Changing these should cover a broad range of models, but this class is not intended to cover all possible use cases. For full flexibility, inherit from tfrs.models.Model and provide your own implementations of the compute_loss and call methods.

embedding_layer The embedding layer is applied to categorical features. It expects a string-to-tensor (or SparseTensor/RaggedTensor) dict as an input, and outputs a dictionary of string-to-tensor of feature_name, embedded_value pairs. {feature_name_i: tensor_i} -> {feature_name_i: emb(tensor_i)}.
bottom_stack The bottom_stack layer is applied to dense features before feature interaction. If None, an MLP with layer sizes [256, 64, 16] is used. For DLRM model, the output of bottom_stack should be of shape (batch_size, embedding dimension).
feature_interaction Feature interaction layer is applied to the bottom_stack output and sparse feature embeddings. If it is None, DotInteraction layer is used.
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.
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.

dense_trainable_variables Returns all trainable variables that are not embeddings.
embedding_trainable_variables Returns trainable variables from embedding tables.

When training a recommendation model with embedding tables, sometimes it's preferable to use separate optimizers/learning rates for embedding variables and dense variables. tfrs.experimental.optimizers.CompositeOptimizer can be used to apply different optimizers to embedding variables and the remaining variables.

Methods

call

View source

call(
    inputs: Dict[str, tf.Tensor]
) -> tf.Tensor

Executes forward and backward pass, returns loss.

Args
inputs Model function inputs (features and labels).

Returns
loss Scalar tensor.