TensorFlow 推荐系统
TensorFlow Recommenders 是一个使用 TensorFlow 构建推荐系统模型的库。
它帮助推荐系统的整个构建流程:数据准备、模型制定、训练、评估和部署。
它基于 Keras 构建,旨在提供轻松的学习曲线,同时给予你构建复杂模型的灵活性。
安装
请确保你已经安装了 TensorFlow 2.x,然后通过 pip
进行安装:
pip install tensorflow-recommenders
文档
快速开始
为 Movielens 100K 数据集构建一个分解模型非常简单 (Colab):
from typing import Dict, Text
import tensorflow as tf
import tensorflow_datasets as tfds
import tensorflow_recommenders as tfrs
# 评分数据。
ratings = tfds.load('movielens/100k-ratings', split="train")
# 所有可用电影的特征。
movies = tfds.load('movielens/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"]))
# 构建模型。
class Model(tfrs.Model):
def __init__(self):
super().__init__()
# 设置用户表示。
self.user_model = tf.keras.layers.Embedding(
input_dim=2000, output_dim=64)
# 设置电影表示。
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))
# 随机打散数据并划分为训练和测试集。
tf.random.set_seed(42)
shuffled = ratings.shuffle(100_000, seed=42, reshuffle_each_iteration=False)
train = shuffled.take(80_000)
test = shuffled.skip(80_000).take(20_000)
# 训练。
model.fit(train.batch(4096), epochs=5)
# 评估。
model.evaluate(test.batch(4096), return_dict=True)