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

2021-06-10 15:28:49 浏览数 (1)


TensorFlow Recommenders

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

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

代码语言:javascript复制
pip install tensorflow-recommenders

TFRS使得:

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

TFRS模块:

datasets:数据集模块

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

layers:图层模块

losses:损失函数模块

metrics:指标模块

models:模型模块

tasks:任务库模块

TFRS例子:

导入库

代码语言:javascript复制
import tensorflow_datasets as tfds
import tensorflow_recommenders as tfrs

导入数据

代码语言:javascript复制
# 评分数据.
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"]))

构建模型

代码语言:javascript复制
# 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)
代码语言:javascript复制
    # 在整个候选数据集上设置检索任务和评估指标。
    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)

模型训练和评估

代码语言:javascript复制
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

0 人点赞