TensorFlow Recommenders 现已开源,让推荐系统更上一层楼!

文 / Maciej Kula 和 James Chen,Google Brain

推荐系统是机器学习的一大主要应用,能够根据用户偏好推送相关内容,比如推荐电影、餐厅、搭配时装首饰或筛选媒体信息流等。

Google 过去几年一直在探索新的深度学习技术,力求通过结合多任务学习、强化学习、提取 更好的用户表征 和构建 公平性指标 提供更好的推荐。这些努力和其他方面的进展大幅改善了我们的推荐效果。

今天,我们荣幸地推出 TensorFlow Recommenders (TFRS),这款开源 TensorFlow 软件包可简化构建、评估和应用复杂的推荐模型。

TFRS 使用 TensorFlow 2.x 构建,有助于:

TFRS 基于 TensorFlow 2.x 和 Keras,十分易于上手,在采用模块化设计的同时(您可以自定义每个层和评价指标),仍然组成了一个强有力的整体(各个组件可以良好协作)。在 TFRS 的设计过程中,我们一直强调灵活性和易用性:合理的默认设置、直观易行的常见任务以及更复杂或自定义的推荐任务。

TensorFlow Recommenders 现已在 GitHub 上开源。我们的目标是让其不断发展,能够灵活地进行学术研究,并以高度可扩展的方式构建全网推荐系统。我们还计划在多任务学习、特征交叉建模、自监督学习和最前沿 (SOTA) 近似最邻近计算 方面扩展其功能。

示例:构建电影推荐工具

让我们先用一个简单示例展现 TensorFlow Recommenders 的使用方法。首先,使用 pip 安装 TFRS:

!pip install tensorflow_recommenders

然后,我们可以使用 MovieLens 数据集 训练一个简单的电影推荐模型。数据集所含信息包括用户观看了哪些电影以及用户对该电影的评分。

我们将使用这一数据集构建模型,预测用户已观看和未观看的电影。此类任务通常选择双塔模型:一个具有两个子模型的神经网络,分别学习 query 和 candidate 的表征。给定的 query-candidate 对 的得分 (score) 只是这两个塔的输出的点积。

这个模型架构相当灵活。query 塔的输入可以是:用户 ID、搜索关键词或时间戳;对于 candidate 侧则有:电影片名、描述、梗概、主演名单。

在此示例中,我们在 query 塔仅使用用户 ID,在 candidate 塔仅使用电影片名。

我们先来准备数据。数据可从 TensorFlow Datasets 获取。

import tensorflow as tf
import tensorflow_datasets as tfds
import tensorflow_recommenders as tfrs
# Ratings data.
ratings = tfds.load("movie_lens/100k-ratings", split="train")
# Features of all the available movies.
movies = tfds.load("movie_lens/100k-movies", split="train")

在数据集的所有可用特征中,最实用的是用户 ID 和电影片名。虽然 TFRS 有多种可选特征,但为简单起见,我们只使用这两项。

ratings = ratings.map(lambda x: {
    "movie_title": x["movie_title"],
    "user_id": x["user_id"],
})
movies = movies.map(lambda x: x["movie_title"])

只使用用户 ID 和电影片名时,我们简单的 双塔模型 与典型的矩阵分解模型非常相似。我们需要使用以下内容进行构建:

  • 一个用户塔,将用户 ID 转换为用户 embedding 向量(高维向量表示)。
  • 一个电影塔,将电影片名转换为电影 embedding 向量。
  • 一个损失函数,对于观看行为,最大化预测用户与电影的匹配度,而未观看的行为进行最小化。

TFRS 和 Keras 为实现这一目标提供了大量基本模块。我们可以从创建模型类开始。在__init__ 方法中,我们设置一些超参数以及模型的主要组件。

class TwoTowerMovielensModel(tfrs.Model):
  """MovieLens prediction model."""

  def __init__(self):
    # The `__init__` method sets up the model architecture.
    super().__init__()

    # How large the representation vectors are for inputs: larger vectors make
    # for a more expressive model but may cause over-fitting.
    embedding_dim = 32
    num_unique_users = 1000
    num_unique_movies = 1700
    eval_batch_size = 128

第一个主要组件是用户模型:一组描述如何将原始用户特征转换为数字化用户表征的层。我们在这里使用 Keras 预处理层 将用户 ID 转换为整数索引,然后将其映射到学习的 embedding 向量:

# Set up user and movie representations.
    self.user_model = tf.keras.Sequential([
      # We first turn the raw user ids into contiguous integers by looking them
      # up in a vocabulary.
      tf.keras.layers.experimental.preprocessing.StringLookup(
          max_tokens=num_unique_users),
      # We then map the result into embedding vectors.
      tf.keras.layers.Embedding(num_unique_users, embedding_dim)
    ])

电影模型看起来很相似,能够将电影片名转换为 embedding 向量:

self.movie_model = tf.keras.Sequential([
      tf.keras.layers.experimental.preprocessing.StringLookup(
          max_tokens=num_unique_movies),
      tf.keras.layers.Embedding(num_unique_movies, embedding_dim)
    ])

得到用户和电影模型后,就需要定义我们的目标和它的评估指标了。在 TFRS 中,可以通过 Retrieval 任务完成这一点(使用 in-batch softmax loss):

# The `Task` objects has two purposes: (1) it computes the loss and (2)
    # keeps track of metrics.
    self.task = tfrs.tasks.Retrieval(
        # In this case, our metrics are top-k metrics: given a user and a known
        # watched movie, how highly would the model rank the true movie out of
        # all possible movies?
        metrics=tfrs.metrics.FactorizedTopK(
            candidates=movies.batch(eval_batch_size).map(self.movie_model)
        )
    )

我们使用 compute_loss 方法查看模型的训练过程:

def compute_loss(self, features, training=False):
    # The `compute_loss` method determines how loss is computed.

    # Compute user and item embeddings.
    user_embeddings = self.user_model(features["user_id"])
    movie_embeddings = self.movie_model(features["movie_title"])

    # Pass them into the task to get the resulting loss. The lower the loss is, the
    # better the model is at telling apart true watches from watches that did
    # not happen in the training data.
    return self.task(user_embeddings, movie_embeddings)

我们可以调用 Keras 的 fit 拟合此模型:

model = MovielensModel()
model.compile(optimizer=tf.keras.optimizers.Adagrad(0.1))

model.fit(ratings.batch(4096), verbose=False)

要对模型的推荐进行 Sanity-Check(合理性检验),我们可以使用 TFRS BruteForce 层。BruteForce 层以预先计算好的 candidate 的表征进行排序,允许我们对所有可能的 candidate 计算其所在 query-candidate 对的得分,并返回排名最靠前的电影 (query):

index = tfrs.layers.ann.BruteForce(model.user_model)
index.index(movies.batch(100).map(model.movie_model), movies)

# Get recommendations.
_, titles = index(tf.constant(["42"]))
print(f"Recommendations for user 42: {titles[0, :3]}")

当然,BruteForce 层只适用于非常小的数据集。有关将 TFRS 与近似最邻近库 Annoy 结合使用的示例,请参阅我们的 完整教程

我们希望这能让您对 TensorFlow Recommenders 的功能有所了解。要了解更多信息,请查看我们的 教程API 参考。如果您想参与,一同推动 TensorFlow 推荐系统发展,请考虑 贡献您的一份力量!我们还将在近期宣布成立 TensorFlow Recommendations 特殊兴趣小组 (SIG),欢迎大家就嵌入向量学习和分布式训练与应用等主题开展合作和做出贡献。敬请期待!

致谢

TensorFlow Recommenders 是 Google 以及其他组织的人员共同努力的成果。我们要感谢 Tiansheng Yao、Xinyang Yi、Ji Yang 对库的核心贡献,感谢 Lichan Hong 和 Ed Chi 的领导与指导。我们也要感谢 Zhe Zhao、Derek Cheng、Sagar Jain、Alexandre Passos、Francois Chollet、Sandeep Gupta、Eric Ni 等人对项目的建议和支持。

image

原文:Introducing TensorFlow Recommenders
中文:TensorFlow 公众号