TensorFlow新手实操：如何在TensorFlow Recommenders库中使用自定义CSV数据集替代Movielens训练推荐模型

嘿，我来帮你一步步搞定怎么把TensorFlow Recommenders（TFRS）示例里的Movielens数据集换成你自己的CSV数据，刚好你的TensorFlow版本我也熟，放心跟着做就行～

你的CSV结构（user_id/item_id/rating）完美适配推荐模型的训练需求，下面是具体步骤：

1. 加载自定义CSV数据集 #

我们可以先用Pandas读取CSV，再转换成TensorFlow的tf.data.Dataset，这样处理起来更灵活：

import pandas as pd
import tensorflow as tf

# 替换成你的CSV文件路径
df = pd.read_csv("your_custom_data.csv")

# 转换成TF数据集，注意对齐数据类型
dataset = tf.data.Dataset.from_tensor_slices({
    "user_id": tf.cast(df["user_id"].values, tf.int32),
    "item_id": tf.cast(df["item_id"].values, tf.int32),
    "rating": tf.cast(df["rating"].values, tf.float32)
})

如果你的CSV没有表头（第一行就是数据），记得加上header=None并指定列名：

df = pd.read_csv("your_custom_data.csv", names=["user_id", "item_id", "rating"])

2. 预处理数据集（和示例逻辑对齐） #

TFRS示例里会对Movielens做打乱、分批等操作，我们对自己的数据集做同样处理即可：

# 打乱数据集，buffer_size建议设为比数据集总条数稍大的值
shuffled_dataset = dataset.shuffle(10000).batch(64)

# 划分训练集和测试集（比如按8:2比例）
train_size = int(0.8 * len(df))
train_dataset = shuffled_dataset.take(train_size)
test_dataset = shuffled_dataset.skip(train_size)

这里的buffer_size和batch_size可以根据你的数据集大小调整，示例里的参数直接复用也没问题。

3. 替换示例中的Movielens加载代码 #

打开TFRS的示例代码（比如入门的MovieLens推荐示例），找到加载Movielens的部分，一般是这样的：

# 示例中原加载Movielens的代码
ratings = tfds.load("movielens/100k-ratings", split="train")
movies = tfds.load("movielens/100k-movies", split="train")

直接把这部分替换成我们自己的数据集。另外，示例里可能用到电影的额外特征（比如标题），但你的数据集只有item_id，所以我们可以简化模型，只使用user_id和item_id作为特征：

首先提取数据集中的唯一用户ID和物品ID，用于嵌入层的词汇表：

# 获取所有唯一的user_id和item_id
unique_user_ids = df["user_id"].unique()
unique_item_ids = df["item_id"].unique()

然后构建用户和物品模型：

from tensorflow_recommenders import layers

# 用户模型：输入user_id，输出嵌入向量
user_model = tf.keras.Sequential([
    layers.StringLookup(vocabulary=unique_user_ids, mask_token=None),
    layers.Embedding(len(unique_user_ids) + 1, 64)
])

# 物品模型：输入item_id，输出嵌入向量
item_model = tf.keras.Sequential([
    layers.StringLookup(vocabulary=unique_item_ids, mask_token=None),
    layers.Embedding(len(unique_item_ids) + 1, 64)
])

如果你的user_id/item_id是纯整数，也可以用IntegerLookup替代StringLookup，在你的TF2.4版本里两者都支持，效果一致。

4. 模型训练（直接复用示例逻辑） #

因为我们的数据集结构和Movielens的ratings数据集完全对齐，所以示例里的训练逻辑几乎不用改。比如构建一个评分预测（Ranking）模型：

from tensorflow_recommenders.models import Ranking

# 构建Ranking模型
ranking_model = Ranking(
    user_model=user_model,
    item_model=item_model,
    task=tf.keras.losses.MeanSquaredError()
)

# 编译模型，优化器和学习率可以沿用示例参数
ranking_model.compile(optimizer=tf.keras.optimizers.Adagrad(learning_rate=0.1))

# 开始训练
history = ranking_model.fit(train_dataset, epochs=5, validation_data=test_dataset)

如果是做召回（Retrieval）模型，只需要把模型换成Retrieval，任务改成RetrievalTask即可，逻辑完全一致。

小提示 #

• 检查数据类型：如果你的user_id/item_id是字符串类型，记得把tf.cast改成tf.string，StringLookup会自动处理。
• 你的版本组合（tensorflow2.4.0、tensorflow-recommenders0.4.0）完全兼容上述代码，不用担心版本冲突。

内容的提问来源于stack exchange，提问作者Amir Masoud