1. 核心思路
FM核心算法不动，详见：基于FM的CTR预测.
只需要更改数据加载函数和训练函数，不过为了方便，FM的实现也会贴在这里。
训练过程中使用的数据集是'ml-100k':
评分数据在 u.data 文件中，一共4列，分别是： user_id, item_id, rating, timestep, 列之间使用制表符(\t)分隔。 但是注意：u.data中并没有特征名称，第一行就是数据。
所以从特征数量上考虑，MoiveLens要比criteo好处理很多。
我们先处理数据加载，然后更改训练函数，最后把代码核心FM贴上来。
2. 代码实现
a. 数据加载
先把完整代码放上来，然后一点一点解释。
def load_movielens_100k(data_path='.', test_size=0.2, random_state=42, threshold=3): # 1. 读取 u.data data = pd.read_csv( f'{data_path}/u.data', sep='\t', header=None, names=['user_id', 'item_id', 'rating', 'timestamp'] ) # 2. 将评分转为二分类标签 data['label'] = (data['rating'] >= threshold).astype(int) # 3. 获取用户数和电影数（ID从1开始，因此最大值就是数量） n_users = data['user_id'].max() # 943 n_items = data['item_id'].max() # 1682 n = n_users + n_items # 2625 # 4. 构建独热编码矩阵（稠密矩阵，2625列，内存可接受） user_ids = data['user_id'].values - 1 # 转为0-based索引 item_ids = data['item_id'].values - 1 X = np.zeros((len(data), n), dtype=np.float32) X[np.arange(len(data)), user_ids] = 1.0 X[np.arange(len(data)), n_users + item_ids] = 1.0 y = data['label'].values.astype(np.float32) # 5. 划分训练集和测试集 X_train, X_test, y_train, y_test = train_test_split( X, y, test_size=test_size, random_state=random_state ) # 6. 转为 PyTorch Tensor X_train = torch.tensor(X_train, dtype=torch.float32) y_train = torch.tensor(y_train, dtype=torch.float32).view(-1, 1) X_test = torch.tensor(X_test, dtype=torch.float32) y_test = torch.tensor(y_test, dtype=torch.float32).view(-1, 1) print(f" 用户数: {n_users}, 电影数: {n_items}") print(f" 特征总数: {n}") print(f" 总样本数: {len(data)}") print(f" 正样本比例: {data['label'].mean():.2%}") print(f" 训练集: {X_train.shape[0]}, 测试集: {X_test.shape[0]}") return (X_train, y_train), (X_test, y_test), n参数
分别表示文件路径、测试集大小、数据划分种子、阈值。
因为总数据量是10000条，所以测试集划分0.2是足够的。
def load_movielens_100k(data_path='.', test_size=0.2, random_state=42, threshold=3):这里面设置阈值(threshold)有3个原因：
- 为了复用CTR的代码套路。(CTR是二分类问题)
- 在经验中验证了阈值和模型性能是有关的，所以为了能方便地调整模型，我们就使用这个参数。
- 如果直接使用评分的数值，模型可能更不健壮、泛用性更差。因为我们更关心的是用户”喜欢“或者”不喜欢“这样的二元状态，而非”评分是多少“这样的回归问题。
读取数据到data
data = pd.read_csv( f'{data_path}/u.data', # 文件路径 sep='\t', # 分隔符 header=None, # 没有列名 names=['user_id', 'item_id', 'rating', 'timestamp'] # 手动定义列名 )将label列从1->5的数值评分类型转换成布尔值
data['label'] = (data['rating'] >= threshold).astype(int)特征独热化
我们先不考虑时间戳，所以一共就两个特征：user和item。
我们决定自己创建独热矩阵(原因等下再说)。为此，要先统计user特征的unique值 和item特征的unique值。好在，当前数据集中user和item都是从1开始的连续正整数值，所以只需用max()函数就可以。
如果想要更具一般性，可以使用unique()，nunique()或者LabelEncoder。
n_users = data['user_id'].max() # 943 n_items = data['item_id'].max() # 1682 n = n_users + n_items # 2625创建独热矩阵：
user_ids = data['user_id'].values - 1 item_ids = data['item_id'].values - 1 X = np.zeros((len(data), n), dtype=np.float32) X[np.arange(len(data)), user_ids] = 1.0 X[np.arange(len(data)), n_users + item_ids] = 1.0显然 X 是独热矩阵，列数为n——两个特征的unique值之和。
Q1.那么行数为什么不是user数呢？ 在这个数据集中，一个样本就是“用户 A 对电影 B 给出评分 C”这一条记录，而用户A显然是可以看很多的电影，从而产生很多条记录的。所以独热矩阵的行数是样本数len(data)。
Q2.我们刚刚还遗留了一个问题：为什么要自己创建独热矩阵？可不可以用pd.get_dummies()？
其实这里是完全可以使用get_dummies()函数的。
常见的不使用pd.get_dummies()的说法是：其无法处理特征中未见过的值。但事实上，在手动独热化的过程中也不能实现处理未见过的值（这里为了方便，在划分数据之前就是用了全局最大值，但这并没有造成造成了数据泄露，因为我们只用这个最大值来编码。不过前两篇博客里，在criteo数据集上确实有数据泄露）。
那么既然功能一样为什么还要手动创建呢？
在FM的论文里提到：FM的优势是能狗结合上下文信息。也就是说，之后优化模型的时候还要增加新的特征，为了对数据有绝对的把控，我们这里使用自己动手实现。
分离标签
y = data['label'].values.astype(np.float32)划分数据集
X_train, X_test, y_train, y_test = train_test_split( X, y, test_size=test_size, random_state=random_state )转为张量(Tensor)
X_train = torch.tensor(X_train, dtype=torch.float32) y_train = torch.tensor(y_train, dtype=torch.float32).view(-1, 1) X_test = torch.tensor(X_test, dtype=torch.float32) y_test = torch.tensor(y_test, dtype=torch.float32).view(-1, 1)b. 设置种子函数
import torch import random import numpy as np def set_seed(seed=42): """ 固定所有随机种子，确保实验可复现 """ # 1. Python 内置 random 模块 random.seed(seed) # 2. NumPy 随机数生成器 np.random.seed(seed) # 3. PyTorch CPU 随机数生成器 torch.manual_seed(seed) # 4. PyTorch GPU 随机数生成器（如果使用 GPU） if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) # 确保 CUDA 运算完全可复现（会牺牲部分性能） torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = Falsec. 模型训练
和criteo数据集相比，只需要改动：数据集读取、超参数设置、优化器改为Adam、以及增加AUC评分。
Adam在小规模数据集上收敛更稳定。
准确率在这里会依赖于阈值——当阈值变化，正样本数变化，准确率随之变化。
而AUC是业界常用标准。
from sklearn.metrics import roc_auc_score from sklearn.metrics import accuracy_score def train(file_path, test_size = 0.3): set_seed(42) # MovieLens数据集 (x_train, y_train), (x_test, y_test), n = load_movielens_100k('ml-100k/') k = 16 # 隐向量维度 # criteo数据集 # (x_train, y_train), (x_test, y_test) = read_criteo(file_path, test_size = test_size) # n = x_train.shape[1] # k = 8 # 模型 model = FMs( n = n, k = k, w_reg = 1e-4, v_reg = 1e-4 ) criterion = nn.BCELoss() # optimizer = torch.optim.SGD(model.parameters(), lr=0.0001) optimizer = torch.optim.Adam(model.parameters(), lr=0.05) # 训练 losses = [] epochs = 100 for epoch in range(epochs): y_pred = model(x_train) loss = criterion(y_pred, y_train) optimizer.zero_grad() loss.backward() optimizer.step() losses.append(loss.item()) # 记录损失 print(f'Epoch {epoch}: Loss {loss.item():.4f}') with torch.no_grad(): y_pred = model(x_test) auc = roc_auc_score(y_test.numpy(), y_pred.numpy()) print(f"AUC: {auc:.4f}") # 评估 with torch.no_grad(): y_test_pred = model(x_test) y_pred_binary = (y_test_pred >= 0.5).float().numpy().flatten() y_test_np = y_test.numpy().flatten() acc = accuracy_score(y_pred_binary, y_test_np) print(f'Accuracy {acc:.4f}') # 作图 plt.plot(range(1, 101), losses) plt.xlabel('Epoch') plt.ylabel('Loss') plt.title('Training Loss Curve') plt.grid(True) plt.show()d. FM实现
import torch.nn as nn class FMs(nn.Module): def __init__(self, n, k, w_reg=0.0, v_reg=0.0): """ n: 输入特征的总维度（即特征向量的长度） 例如 Criteo 数据集中为 13（稠密）+ 26（稀疏）= 39 k: 隐向量 vi 的维度，超参数。 对应 FM 公式中二阶交互项的因子分解维度。 w_reg: 一阶线性部分（w）的 L2 正则化系数 v_reg: 二阶交互部分（v）的 L2 正则化系数 """ super(FMs, self).__init__() self.n = n self.k = k # 偏置项（全局偏差） self.w0 = nn.Parameter(torch.zeros(1)) # 一阶特征权重（对应公式中的 w_i） self.w = nn.Parameter(torch.randn(n, 1)) # 二阶交互的隐向量（对应公式中的 v_i） self.v = nn.Parameter(torch.randn(n, k)) # 初始化 nn.init.normal_(self.w, mean=0, std=0.1) nn.init.normal_(self.v, mean=0, std=0.1) self.w_reg = w_reg self.v_reg = v_reg def forward(self, x): if x.dim() != 2: raise ValueError('Input tensor needs to be a 2D tensor') # 线性部分：w0 + Σ w_i * x_i linear_part = self.w0 + torch.mm(x, self.w) # 二阶交互部分（利用公式简化计算，避免 O(n^2) 复杂度） # 公式：1/2 * Σ_f ( (Σ_i v_i,f x_i)^2 - Σ_i (v_i,f x_i)^2 ) inter_part1 = torch.mm(x, self.v) ** 2 inter_part2 = torch.mm(x ** 2, self.v ** 2) inter_part = 0.5 * torch.sum(inter_part1 - inter_part2, dim=1, keepdim=True) # 最终预测（logit + sigmoid） y = linear_part + inter_part return torch.sigmoid(y) def regularization(self): # L2 正则化项（不含偏置 w0，因为偏置通常不需要正则化） w_p = self.w_reg * torch.sum(self.w ** 2) v_p = self.v_reg * torch.sum(self.v ** 2) return w_p + v_pe. 运行
if __name__ == '__main__': train(file_path = 'train.txt')3. 运行结果
Loss曲线：
4. FM结合上下文
要实现论文中所说的结合上下文的功能——加入了用户年龄、性别、职业和电影类型——只需要修改数据加载部分，另外也增加了隐向量长度 k = 32。
这里独热化使用了LabelEncoder。
对年龄使用了全局归一化，这导致了数据泄露。
from sklearn.preprocessing import LabelEncoder, StandardScaler def load_movielens_with_context(data_path='.', test_size=0.2, random_state=42, threshold=3): # 1. 读取评分数据（u.data） ratings = pd.read_csv( f'{data_path}/u.data', sep='\t', header=None, names=['user_id', 'item_id', 'rating', 'timestamp'] ) # 2. 读取用户数据（u.user） users = pd.read_csv( f'{data_path}/u.user', sep='|', header=None, names=['user_id', 'age', 'gender', 'occupation', 'zip_code'] ) # 3. 读取电影数据（u.item） # 注意：u.item 有 24 列，前5列是电影ID、标题、上映日期、视频发布日期、IMDb链接，后面19列是类型标签 item_columns = ['item_id', 'title', 'release_date', 'video_release_date', 'IMDb_url'] + [f'genre_{i}' for i in range(19)] items = pd.read_csv( f'{data_path}/u.item', sep='|', header=None, encoding='latin-1', # 有些电影名包含特殊字符 names=item_columns ) # 4. 合并数据（评分 + 用户特征 + 电影特征） data = ratings.merge(users, on='user_id').merge(items, on='item_id') # 5. 处理标签（二分类） data['label'] = (data['rating'] >= threshold).astype(int) # 6. 对用户ID和电影ID进行 LabelEncoder（用于one-hot索引） user_encoder = LabelEncoder() item_encoder = LabelEncoder() data['user_encoded'] = user_encoder.fit_transform(data['user_id']) data['item_encoded'] = item_encoder.fit_transform(data['item_id']) # 7. 处理类别型上下文特征（用 LabelEncoder） gender_encoder = LabelEncoder() occ_encoder = LabelEncoder() data['gender_encoded'] = gender_encoder.fit_transform(data['gender']) data['occupation_encoded'] = occ_encoder.fit_transform(data['occupation']) # 8. 处理数值型上下文特征（归一化到0~1） scaler = StandardScaler() data['age_norm'] = scaler.fit_transform(data[['age']]) # 注：年龄范围约7~73，归一化后可提高收敛速度 # 9. 处理电影类型（19个类型标签，每个取值为0或1，本身就是one-hot） genre_cols = [f'genre_{i}' for i in range(19)] # 这19列已经是0/1，直接使用 # 10. 构建特征矩阵 X # 特征组成：用户ID one-hot + 电影ID one-hot + 年龄(数值) + 性别编码 + 职业编码 + 19个类型标签 n_users = len(user_encoder.classes_) # 943 n_items = len(item_encoder.classes_) # 1682 n_genders = len(gender_encoder.classes_) # 2 n_occupations = len(occ_encoder.classes_) # 21 # 先创建独热编码部分（用户ID + 电影ID） user_ids = data['user_encoded'].values item_ids = data['item_encoded'].values n_onehot = n_users + n_items X_onehot = np.zeros((len(data), n_onehot), dtype=np.float32) X_onehot[np.arange(len(data)), user_ids] = 1.0 X_onehot[np.arange(len(data)), n_users + item_ids] = 1.0 # 再添加其他上下文特征 X_context = np.column_stack([ data['age_norm'].values.astype(np.float32), # 1列，归一化后的年龄 data['gender_encoded'].values.astype(np.float32), # 1列，性别编码 data['occupation_encoded'].values.astype(np.float32), # 1列，职业编码 data[genre_cols].values.astype(np.float32) # 19列，电影类型 ]) # 拼接所有特征 X = np.concatenate([X_onehot, X_context], axis=1) # 注意：n 现在是一阶和二阶特征的总数 n = X.shape[1] y = data['label'].values.astype(np.float32) # 划分训练集和测试集 X_train, X_test, y_train, y_test = train_test_split( X, y, test_size=test_size, random_state=random_state ) # 转为 PyTorch Tensor X_train = torch.tensor(X_train, dtype=torch.float32) y_train = torch.tensor(y_train, dtype=torch.float32).view(-1, 1) X_test = torch.tensor(X_test, dtype=torch.float32) y_test = torch.tensor(y_test, dtype=torch.float32).view(-1, 1) print(f" 用户数: {n_users}, 电影数: {n_items}") print(f" 特征总数: {n}") print(f" 总样本数: {len(data)}") print(f" 正样本比例: {data['label'].mean():.2%}") print(f" 训练集: {X_train.shape[0]}, 测试集: {X_test.shape[0]}") return (X_train, y_train), (X_test, y_test), n5. 运行结果
Loss曲线：
显然，对比不加上下文的FM，模型的AUC提升了0.08，说明上下文特征确实有用。