Watch keynotes, product sessions, workshops, and more from Google I/O See playlist


A configurable ranking model.

Inherits From: Model

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([

    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 optimizer to embedding variables and the remaining variables.



View source

Executes forward and backward pass, returns loss.

inputs Model function inputs (features and labels).

loss Scalar tensor.


View source

Computes the loss and metrics of the model.

inputs A data structure of tensors of the following format: ({"dense_features": dense_tensor, "sparse_features": sparse_tensors}, label_tensor), or ({"dense_features": dense_tensor, "sparse_features": sparse_tensors}, label_tensor, sample_weight tensor).
training Whether the model is in training mode.

Loss tensor.

ValueError if the the shape of the inputs is invalid.