► KerasRS / 示例 / 电影推荐：排序

电影推荐：排序

作者： Fabien Hertschuh, Abheesht Sharma
创建日期 2025/04/28
最后修改 2025/04/28
描述： 使用双塔模型对电影进行排序。

在 Colab 中查看 • GitHub 源代码

简介

推荐系统通常由两个阶段组成

检索阶段负责从所有可能的候选中选择最初的数百个候选项。该模型的主要目标是有效地筛选掉用户不感兴趣的所有候选项。由于检索模型可能需要处理数百万个候选项，因此它必须计算高效。
排序阶段接收检索模型的输出，并对其进行微调，以选择最佳的少量推荐结果。其任务是将用户可能感兴趣的项目集缩小到可能候选项的短名单。

在本教程中，我们将重点关注第二阶段，即排序。如果您对检索阶段感兴趣，请查看我们的检索教程。

在本教程中，我们将

获取数据并将其分割为训练集和测试集。
实现一个排序模型。
拟合和评估模型。
测试使用模型进行预测。

首先，选择我们希望在其上运行的后端 JAX，并导入所有必要的库。

import os

os.environ["KERAS_BACKEND"] = "jax"  # `"tensorflow"`/`"torch"`

import keras
import tensorflow as tf  # Needed for the dataset
import tensorflow_datasets as tfds

准备数据集

我们将使用与检索教程相同的数据。评分是我们试图预测的目标。

# Ratings data.
ratings = tfds.load("movielens/100k-ratings", split="train")
# Features of all the available movies.
movies = tfds.load("movielens/100k-movies", split="train")

在 Movielens 数据集中，用户 ID 是从 1 开始且没有间隙的整数（表示为字符串）。通常，您需要创建一个查找表将用户 ID 映射到从 0 到 N-1 的整数。但作为简化，我们将直接使用用户 ID 作为模型中的索引，特别是用于从用户嵌入表中查找用户嵌入。因此，我们需要知道用户数量。

users_count = (
    ratings.map(lambda x: tf.strings.to_number(x["user_id"], out_type=tf.int32))
    .reduce(tf.constant(0, tf.int32), tf.maximum)
    .numpy()
)

在 Movielens 数据集中，电影 ID 是从 1 开始且没有间隙的整数（表示为字符串）。通常，您需要创建一个查找表将电影 ID 映射到从 0 到 N-1 的整数。但作为简化，我们将直接使用电影 ID 作为模型中的索引，特别是用于从电影嵌入表中查找电影嵌入。因此，我们需要知道电影数量。

movies_count = movies.cardinality().numpy()

模型的输入是用户 ID 和电影 ID，标签是评分。

def preprocess_rating(x):
    return (
        # Inputs are user IDs and movie IDs
        {
            "user_id": tf.strings.to_number(x["user_id"], out_type=tf.int32),
            "movie_id": tf.strings.to_number(x["movie_id"], out_type=tf.int32),
        },
        # Labels are ratings between 0 and 1.
        (x["user_rating"] - 1.0) / 4.0,
    )

我们将把 80% 的评分放入训练集，20% 放入测试集来分割数据。

shuffled_ratings = ratings.map(preprocess_rating).shuffle(
    100_000, seed=42, reshuffle_each_iteration=False
)
train_ratings = shuffled_ratings.take(80_000).batch(1000).cache()
test_ratings = shuffled_ratings.skip(80_000).take(20_000).batch(1000).cache()

实现模型

架构

排序模型不像检索模型那样面临相同的效率限制，因此我们在架构选择上有一些更大的自由度。

由多个堆叠的全连接层组成的模型是排序任务中相对常见的架构。我们可以按如下方式实现它

class RankingModel(keras.Model):
    """Create the ranking model with the provided parameters.

    Args:
      num_users: Number of entries in the user embedding table.
      num_candidates: Number of entries in the candidate embedding table.
      embedding_dimension: Output dimension for user and movie embedding tables.
    """

    def __init__(
        self,
        num_users,
        num_candidates,
        embedding_dimension=32,
        **kwargs,
    ):
        super().__init__(**kwargs)
        # Embedding table for users.
        self.user_embedding = keras.layers.Embedding(num_users, embedding_dimension)
        # Embedding table for candidates.
        self.candidate_embedding = keras.layers.Embedding(
            num_candidates, embedding_dimension
        )
        # Predictions.
        self.ratings = keras.Sequential(
            [
                # Learn multiple dense layers.
                keras.layers.Dense(256, activation="relu"),
                keras.layers.Dense(64, activation="relu"),
                # Make rating predictions in the final layer.
                keras.layers.Dense(1),
            ]
        )

    def call(self, inputs):
        user_id, movie_id = inputs["user_id"], inputs["movie_id"]
        user_embeddings = self.user_embedding(user_id)
        candidate_embeddings = self.candidate_embedding(movie_id)
        return self.ratings(
            keras.ops.concatenate([user_embeddings, candidate_embeddings], axis=1)
        )

我们首先实例化模型。注意，我们在用户和电影的数量上加了+ 1，以考虑到 ID 零没有被使用（ID 从 1 开始），但仍然在嵌入表中占有一行。

model = RankingModel(users_count + 1, movies_count + 1)

损失和指标

下一个组件是用于训练模型的损失。Keras 有多种损失函数可以方便实现。在这种情况下，我们将使用MeanSquaredError损失来预测评分。我们还将查看RootMeanSquaredError指标。

model.compile(
    loss=keras.losses.MeanSquaredError(),
    metrics=[keras.metrics.RootMeanSquaredError()],
    optimizer=keras.optimizers.Adagrad(learning_rate=0.1),
)

拟合和评估

定义模型后，我们可以使用标准的 Keras model.fit() 来训练模型。

model.fit(train_ratings, epochs=5)

Epoch 1/5
80/80 ━━━━━━━━━━━━━━━━━━━━ 4s 10ms/step - loss: 0.1071 - root_mean_squared_error: 0.3218
Epoch 2/5
80/80 ━━━━━━━━━━━━━━━━━━━━ 1s 3ms/step - loss: 0.0769 - root_mean_squared_error: 0.2773
Epoch 3/5
80/80 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - loss: 0.0745 - root_mean_squared_error: 0.2730
Epoch 4/5
80/80 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - loss: 0.0713 - root_mean_squared_error: 0.2670
Epoch 5/5
80/80 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - loss: 0.0682 - root_mean_squared_error: 0.2612

<keras.src.callbacks.history.History at 0x7f69980ee5f0>

随着模型训练，损失正在下降，RMSE 指标正在改善。

最后，我们可以在测试集上评估我们的模型。RMSE 指标越低，我们的模型在预测评分方面就越准确。

model.evaluate(test_ratings, return_dict=True)

20/20 ━━━━━━━━━━━━━━━━━━━━ 3s 12ms/step - loss: 0.0649 - root_mean_squared_error: 0.2548

{'loss': 0.06562447547912598, 'root_mean_squared_error': 0.2561727464199066}

测试排序模型

到目前为止，我们仅通过 ID 处理电影。现在是时候创建一个以电影 ID 为键的映射，以便能够显示电影标题了。

movie_id_to_movie_title = {
    int(x["movie_id"]): x["movie_title"] for x in movies.as_numpy_iterator()
}
movie_id_to_movie_title[0] = ""  # Because id 0 is not in the dataset.

现在我们可以通过计算一组电影的预测值，然后根据这些预测值对电影进行排序来测试排序模型

user_id = 42
movie_ids = [204, 141, 131]
predictions = model.predict(
    {
        "user_id": keras.ops.array([user_id] * len(movie_ids)),
        "movie_id": keras.ops.array(movie_ids),
    }
)
predictions = keras.ops.convert_to_numpy(keras.ops.squeeze(predictions, axis=1))

for movie_id, prediction in zip(movie_ids, predictions):
    print(f"{movie_id_to_movie_title[movie_id]}: {5.0 * prediction:,.2f}")

1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 153ms/step
b'Back to the Future (1985)': 3.53
b'20,000 Leagues Under the Sea (1954)': 3.26
b"Breakfast at Tiffany's (1961)": 3.43