TFRS | 谷歌开源新一代推荐系统库

张小磊

发布于 2021-06-10 15:28:49

9830

发布于 2021-06-10 15:28:49

TensorFlow Recommenders

TensorFlow推荐器是一个使用TensorFlow构建推荐系统模型的库。它有助于构建推荐系统的全部工作流程:数据准备、模型制定、训练、评估和部署。它构建在Keras上，目标是让学习者有一个平缓的学习曲线，同时仍然给你构建复杂模型的灵活性。

安装(确保安装了TensorFlow 2.x)

pip install tensorflow-recommenders

TFRS使得:

建立并评估灵活的推荐检索模型。
自由地将item、user和上下文信息合并到推荐模型中。
联合训练多目标推荐的多任务模型。

TFRS模块：

datasets：数据集模块

examples：示例中使用的功能模块

layers：图层模块

losses：损失函数模块

metrics：指标模块

models：模型模块

tasks：任务库模块

TFRS例子：

导入库

import tensorflow_datasets as tfds
import tensorflow_recommenders as tfrs

导入数据

# 评分数据.
ratings = tfds.load('movie_lens/100k-ratings', split="train")
# 电影特征数据.
movies = tfds.load('movie_lens/100k-movies', split="train")

# 选取特征.
ratings = ratings.map(lambda x: {
    "movie_id": tf.strings.to_number(x["movie_id"]),
    "user_id": tf.strings.to_number(x["user_id"])
})
movies = movies.map(lambda x: tf.strings.to_number(x["movie_id"]))

构建模型

# Build a model.
class Model(tfrs.Model):

  def __init__(self):
    super().__init__()

    # user embedding.
    self.user_model = tf.keras.layers.Embedding(
        input_dim=2000, output_dim=64)
    # movie embedding.
    self.item_model = tf.keras.layers.Embedding(
        input_dim=2000, output_dim=64)

    # 在整个候选数据集上设置检索任务和评估指标。
    self.task = tfrs.tasks.Retrieval(
        metrics=tfrs.metrics.FactorizedTopK(
            candidates=movies.batch(128).map(self.item_model)
        )
    )

  def compute_loss(self, features: Dict[Text, tf.Tensor], training=False) -> tf.Tensor:

    user_embeddings = self.user_model(features["user_id"])
    movie_embeddings = self.item_model(features["movie_id"])

    return self.task(user_embeddings, movie_embeddings)

模型训练和评估

model = Model()
model.compile(optimizer=tf.keras.optimizers.Adagrad(0.5))

# Randomly shuffle data and split between train and test.
tf.random.set_seed(42)
shuffled = ratings.shuffle(100_000, seed=42, reshuffle_each_iteration=False)

train = shuffled.take(80_000) # 取前8000做训练集
test = shuffled.skip(80_000).take(20_000) # 取80000-100000为测试集

# Train.
model.fit(train.batch(4096), epochs=5)

# Evaluate.
model.evaluate(test.batch(4096), return_dict=True)

更多细节请移步：

源码：

https://github.com/tensorflow/recommenders

指南:

https://www.tensorflow.org/recommenders/examples/quickstart

API：

https://www.tensorflow.org/recommenders/api_docs/tfrs

本文参与腾讯云自媒体分享计划，分享自微信公众号。

原始发表：2021-05-28，如有侵权请联系 cloudcommunity@tencent.com 删除

tensorflow