 在 TensorFlow.org 上查看
|
 在 Google Colab 中运行
|
 在 GitHub 上查看源代码
|
 下载笔记本
|
在本教程中，我们使用 TF-Ranking 构建了一个简单的双塔排名模型，使用 MovieLens 100K 数据集。我们可以使用此模型根据预测的用户评分对电影进行排名和推荐。
设置
安装并导入 TF-Ranking 库
pip install -q tensorflow-rankingpip install -q --upgrade tensorflow-datasets
from typing import Dict, Tuple
import tensorflow as tf
import tensorflow_datasets as tfds
import tensorflow_ranking as tfr
2024-03-19 11:34:49.704174: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered 2024-03-19 11:34:49.704225: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered 2024-03-19 11:34:49.705795: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
读取数据
通过创建评分数据集和电影数据集来准备训练模型。使用 user_id 作为查询输入特征，movie_title 作为文档输入特征，以及 user_rating 作为标签来训练排名模型。
%%capture --no-display
# Ratings data.
ratings = tfds.load('movielens/100k-ratings', split="train")
# Features of all the available movies.
movies = tfds.load('movielens/100k-movies', split="train")
# Select the basic features.
ratings = ratings.map(lambda x: {
"movie_title": x["movie_title"],
"user_id": x["user_id"],
"user_rating": x["user_rating"]
})
2024-03-19 11:34:53.385017: E external/local_xla/xla/stream_executor/cuda/cuda_driver.cc:274] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected
构建词汇表，将所有用户 ID 和所有电影标题转换为嵌入层的整数索引
movies = movies.map(lambda x: x["movie_title"])
users = ratings.map(lambda x: x["user_id"])
user_ids_vocabulary = tf.keras.layers.experimental.preprocessing.StringLookup(
mask_token=None)
user_ids_vocabulary.adapt(users.batch(1000))
movie_titles_vocabulary = tf.keras.layers.experimental.preprocessing.StringLookup(
mask_token=None)
movie_titles_vocabulary.adapt(movies.batch(1000))
按 user_id 分组，为排名模型形成列表
key_func = lambda x: user_ids_vocabulary(x["user_id"])
reduce_func = lambda key, dataset: dataset.batch(100)
ds_train = ratings.group_by_window(
key_func=key_func, reduce_func=reduce_func, window_size=100)
for x in ds_train.take(1):
for key, value in x.items():
print(f"Shape of {key}: {value.shape}")
print(f"Example values of {key}: {value[:5].numpy()}")
print()
Shape of movie_title: (100,) Example values of movie_title: [b'Man Who Would Be King, The (1975)' b'Silence of the Lambs, The (1991)' b'Next Karate Kid, The (1994)' b'2001: A Space Odyssey (1968)' b'Usual Suspects, The (1995)'] Shape of user_id: (100,) Example values of user_id: [b'405' b'405' b'405' b'405' b'405'] Shape of user_rating: (100,) Example values of user_rating: [1. 4. 1. 5. 5.]
生成批处理特征和标签
def _features_and_labels(
x: Dict[str, tf.Tensor]) -> Tuple[Dict[str, tf.Tensor], tf.Tensor]:
labels = x.pop("user_rating")
return x, labels
ds_train = ds_train.map(_features_and_labels)
ds_train = ds_train.apply(
tf.data.experimental.dense_to_ragged_batch(batch_size=32))
WARNING:tensorflow:From /tmpfs/tmp/ipykernel_12750/4021484596.py:10: dense_to_ragged_batch (from tensorflow.python.data.experimental.ops.batching) is deprecated and will be removed in a future version. Instructions for updating: Use `tf.data.Dataset.ragged_batch` instead.
在 ds_train 中生成的 user_id 和 movie_title 张量形状为 [32, None]，其中第二维在大多数情况下为 100，除了在列表中分组的项目少于 100 的批次。因此，使用在不规则张量上工作的模型。
for x, label in ds_train.take(1):
for key, value in x.items():
print(f"Shape of {key}: {value.shape}")
print(f"Example values of {key}: {value[:3, :3].numpy()}")
print()
print(f"Shape of label: {label.shape}")
print(f"Example values of label: {label[:3, :3].numpy()}")
Shape of movie_title: (32, None) Example values of movie_title: [[b'Man Who Would Be King, The (1975)' b'Silence of the Lambs, The (1991)' b'Next Karate Kid, The (1994)'] [b'Flower of My Secret, The (Flor de mi secreto, La) (1995)' b'Little Princess, The (1939)' b'Time to Kill, A (1996)'] [b'Kundun (1997)' b'Scream (1996)' b'Power 98 (1995)']] Shape of user_id: (32, None) Example values of user_id: [[b'405' b'405' b'405'] [b'655' b'655' b'655'] [b'13' b'13' b'13']] Shape of label: (32, None) Example values of label: [[1. 4. 1.] [3. 3. 3.] [5. 1. 1.]]
定义模型
通过继承 tf.keras.Model 并实现 call 方法来定义排名模型
class MovieLensRankingModel(tf.keras.Model):
def __init__(self, user_vocab, movie_vocab):
super().__init__()
# Set up user and movie vocabulary and embedding.
self.user_vocab = user_vocab
self.movie_vocab = movie_vocab
self.user_embed = tf.keras.layers.Embedding(user_vocab.vocabulary_size(),
64)
self.movie_embed = tf.keras.layers.Embedding(movie_vocab.vocabulary_size(),
64)
def call(self, features: Dict[str, tf.Tensor]) -> tf.Tensor:
# Define how the ranking scores are computed:
# Take the dot-product of the user embeddings with the movie embeddings.
user_embeddings = self.user_embed(self.user_vocab(features["user_id"]))
movie_embeddings = self.movie_embed(
self.movie_vocab(features["movie_title"]))
return tf.reduce_sum(user_embeddings * movie_embeddings, axis=2)
创建模型，然后使用排名 tfr.keras.losses 和 tfr.keras.metrics 进行编译，它们是 TF-Ranking 包的核心。
此示例使用排名特定的 **softmax 损失**，这是一种列表级损失，旨在提高排名列表中所有相关项目的排名，使其更有可能出现在无关项目的顶部。与多类分类问题中的 softmax 损失（其中只有一个类是正类，其余是负类）不同，TF-Ranking 库支持查询列表中的多个相关文档和非二元相关性标签。
对于排名指标，此示例专门使用 **归一化折现累积增益 (NDCG)** 和 **平均倒数排名 (MRR)**，它们计算排名查询列表的用户效用，并使用位置折扣。有关排名指标的更多详细信息，请查看评估指标 离线指标。
# Create the ranking model, trained with a ranking loss and evaluated with
# ranking metrics.
model = MovieLensRankingModel(user_ids_vocabulary, movie_titles_vocabulary)
optimizer = tf.keras.optimizers.Adagrad(0.5)
loss = tfr.keras.losses.get(
loss=tfr.keras.losses.RankingLossKey.SOFTMAX_LOSS, ragged=True)
eval_metrics = [
tfr.keras.metrics.get(key="ndcg", name="metric/ndcg", ragged=True),
tfr.keras.metrics.get(key="mrr", name="metric/mrr", ragged=True)
]
model.compile(optimizer=optimizer, loss=loss, metrics=eval_metrics)
训练和评估模型
使用 model.fit 训练模型。
model.fit(ds_train, epochs=3)
Epoch 1/3 48/48 [==============================] - 7s 56ms/step - loss: 998.7637 - metric/ndcg: 0.8213 - metric/mrr: 1.0000 Epoch 2/3 48/48 [==============================] - 4s 53ms/step - loss: 997.1824 - metric/ndcg: 0.9161 - metric/mrr: 1.0000 Epoch 3/3 48/48 [==============================] - 4s 53ms/step - loss: 994.8384 - metric/ndcg: 0.9383 - metric/mrr: 1.0000 <keras.src.callbacks.History at 0x7f666424d700>
生成预测并评估。
# Get movie title candidate list.
for movie_titles in movies.batch(2000):
break
# Generate the input for user 42.
inputs = {
"user_id":
tf.expand_dims(tf.repeat("42", repeats=movie_titles.shape[0]), axis=0),
"movie_title":
tf.expand_dims(movie_titles, axis=0)
}
# Get movie recommendations for user 42.
scores = model(inputs)
titles = tfr.utils.sort_by_scores(scores,
[tf.expand_dims(movie_titles, axis=0)])[0]
print(f"Top 5 recommendations for user 42: {titles[0, :5]}")
Top 5 recommendations for user 42: [b'Star Wars (1977)' b'Liar Liar (1997)' b'Toy Story (1995)' b'Raiders of the Lost Ark (1981)' b'Sound of Music, The (1965)']