简介:面向推荐系统算法工程师和CTR入门学习者,这份PyTorch实现DeepFM的代码包完整覆盖了数据预处理、FM与DNN双分支构建、模型训练和效果评估等关键环节,可帮助读者理解特征二阶交互与高阶非线性组合在点击率预测中的联合建模方式。压缩包共37个文件,整体约40KB,以Python脚本、PyTorch模型文件、README说明及项目配置文件为主,结构紧凑,适合直接阅读和二次改造。包内main.py提供训练入口,model/DeepFM.py定义网络结构,utils/dataPreprocess.py完成特征编码与数据集划分,README对整体流程与设计思路做了梳理;对嵌入向量内积计算、多层全连接堆叠和交叉熵损失训练均有可运行实现。目前已有666人学习浏览,适合想快速跑通DeepFM基线并深入源码的推荐系统开发者。 在CTR预估和推荐系统这个圈子里,DeepFM 算是一个绕不开的经典模型了。它把低阶特征交叉和高阶特征交互放在同一个框架里做端到端训练,效果扎实,工程上也好落地。最近我在 PyTorch 里重新实现了一遍 DeepFM,把整个流程跑通之后,对很多细节的体会比看论文时要深得多。这篇博文就从一个实际项目的角度,分享 DeepFM 的原理拆解、PyTorch 实现的关键代码、训练调参的经验,以及那些正式文档里不会写但实际踩过的坑。
1. 为什么 CTR 预估场景里 DeepFM 这么能打
1.1 CTR 预估的本质难点在哪里
CTR 预估(Click-Through Rate Prediction)做的是预测用户点击某个物品的概率。在推荐系统、广告投放这些场景里,输入的特征有一个非常明显的特点:极度稀疏且有大量类别型特征。比如用户 ID、物品 ID、城市 ID、设备类型,这些特征的取值空间动辄百万甚至千万级别,但单个样本里只有少数几个位置是激活的。
面对这种特征空间,传统机器学习算法比如 LR(逻辑回归),虽然简单高效,但表达不了特征之间的非线性关系。实际业务里,用户点击行为往往是由“用户属性 + 物品属性 + 上下文环境”三者的交互共同决定的。比如“男性用户 + 体育类物品”这个组合,比单独看“男性用户”或“体育类物品”的预测价值高得多。如果模型只能学线性结构,就必须人工去组合这些交叉特征,成本很高,而且不可能穷尽所有组合。
1.2 FM 模型的局限和 DeepFM 的出现
FM(Factorization Machine)解决了 LR 不能自动学习二阶特征交叉的问题。它的核心思路是为每个特征学一个隐向量,两个特征交叉的权重通过隐向量的内积计算,这样即使某些特征组合在训练样本中从未出现,也能通过隐向量泛化出对应的交叉权重。
但 FM 只到二阶交叉就停了,要去拟合更高阶的特征交互就必须手动设计或者堆阶数,计算成本和训练难度都会膨胀。另一边,深度学习模型如 FNN、PNN 可以学高阶特征交互,但前者依赖预训练的 FM 初始化权重,训练流程割裂,后者对低阶特征交叉的表达又不充分。
DeepFM 的思路很简单粗暴:用 FM 部分负责一阶和二阶特征交叉,用 Deep 部分负责高阶自动特征提取,两部分共享同一个 Embedding 层,联合训练。这样低阶交叉和高阶交互都不用人工设计,一套模型端到端搞定。线上效果稳定,结构简洁,这也是它在工业界口碑一直很好的原因。
1.3 为什么我选择用 PyTorch 复现
之前用 TensorFlow 跑过这类模型,但最近在 PyTorch 里实现了一遍之后,我对这套架构的理解明显更深了一层。PyTorch 的动态图机制极大方便了调试——网络结构有问题、梯度传播不正常,直接在中间位置加个
print(x.shape)就能快速定位,不需要像静态图那样编译之后再排查。而且 PyTorch 生态里做特征处理的工具也越来越全,写起来更顺手。如果你对这个模型的细节还停留在“大概知道”的程度,自己动手用 PyTorch 撸一遍代码,收益比单纯读论文大得多。
2. 模型核心结构与背后的设计思路
2.1 DeepFM 的几个关键构成模块
DeepFM 的整体结构可以拆成三块:稀疏特征 Embedding 层、FM 部分、Deep 部分。FM 部分里又分成一阶特征线性加权和二阶特征隐向量内积,Deep 部分就是一个多层全连接网络。
Embedding 层是所有稀疏特征进入模型的第一道门槛。每个类别型特征都有一个独立的 Embedding 表,维度通常是 8 到 16。这个维度的选择很讲究——太小了表达不了特征语义,太大了会把参数量推到很夸张的位置。在 Criteo 这种量级的数据上,特征维度动辄几十亿,Embedding 表的大小直接决定模型能不能部署上线。
这里必须强调一个细节:同一个人口统计特征(比如用户 ID)在 FM 二阶部分和 Deep 部分必须共享同一个 Embedding 向量,而不是各学各的。这是 DeepFM 架构设计的核心思想,否则两个部分各自维护一套向量,特征表达就分裂了,低阶和高阶交叉的一致性也会被破坏。
2.2 FM 部分如何做二阶特征交叉
FM 的二阶交叉在数学上可以写成:
y_fm = w0 + sum(w_i * x_i) + sum(sum(<v_i, v_j> * x_i * x_j))直接按这个公式计算的话,复杂度是 O(n^2),在特征维度高的场景下没法用。但 FM 的经典化简技巧可以把复杂度降到 O(nk),方法是把二阶交叉部分的求和顺序换一下:
sum_i sum_j <v_i, v_j> * x_i * x_j = 0.5 * (sum_i (v_i * x_i))^2 - sum_i (v_i * x_i)^2)逻辑是先算所有特征的隐向量加权和,再算整体向量内积,减去每个特征单独平方的求和,本质上就是“平方的求和减求和的平方”。这个优化在工程实现里是标配,不仅速度快,数值也更稳定。下面代码里用的就是这种写法。
2.3 Deep 部分为什么能补上高阶交叉
FM 部分只会显式建模到二阶交叉,而 Deep 部分通过多层非线性变换,理论上可以拟合任意阶数的特征交互。
Deep 部分的输入不能直接丢稀疏 one-hot 向量,而是要把所有特征的 Embedding 拼接成一个稠密向量,再喂给多层全连接网络。每一层的神经元对输入进行线性变换加非线性激活,这种堆叠结构天然具备组合高阶特征的能力。激活函数一般用 ReLU,因为计算快、不容易梯度消失,在大多数 CTR 场景下表现都比 tanh 好。
还需要注意一点:Deep 部分的层数和每层神经元个数决定了模型表达能力的天花板。太浅了学不到高阶交互,太深了容易过拟合、训练变慢。依据我的实践经验,2 到 3 层 128 到 256 维的隐藏层在大多数数据集上性价比最高,继续加深带来的提升很小,但计算开销和过拟合风险都会上来。
3. 基于 PyTorch 的核心代码实现
3.1 特征处理与 Dataset 构建
实现 DeepFM 之前,第一步是把特征处理成模型可以消费的格式。我们假设输入有两类特征:
- 数值特征:比如用户年龄、物品价格,直接做标准化。
- 类别特征:比如用户 ID、物品 ID,做 Label Encoding 之后映射成连续的整数索引。
这里比较常见的坑是:类别特征中出现训练集里没见过的取值。线上推理时,用户或物品的 ID 经常会出现新值,如果你直接查 Embedding 表,可能因为 index 越界而崩溃。所以编码的时候要预留一个未登录索引(比如 0),把所有未知值统一映射到这个索引上。
构造 Dataset 的时候,我习惯把类别特征的索引和数值特征分开存放在两个 list 里,喂给模型时分别处理。这样代码更清晰,也方便后续加特征。
class DeepFMDataset(Dataset): def __init__(self, dense_features, sparse_features, labels): self.dense = torch.tensor(dense_features, dtype=torch.float32) self.sparse = torch.tensor(sparse_features, dtype=torch.long) self.labels = torch.tensor(labels, dtype=torch.float32).unsqueeze(1) def __len__(self): return len(self.labels) def __getitem__(self, idx): return self.dense[idx], self.sparse[idx], self.labels[idx]3.2 Embedding 层的权重共享实现
PyTorch 里对每个稀疏特征维护一个nn.Embedding,然后把这些 Embedding 放进nn.ModuleList里统一管理。每个特征的 Embedding 维度可以独立设置,但在实践中维持一致(比如都设成 8),方便后续拼接。
self.embedding_layers = nn.ModuleList([ nn.Embedding(feat_size, embed_dim, sparse=True) for feat_size in sparse_feat_sizes ])这里有个实际工程问题:如果每个特征都单独搞一个 Embedding 表,参数总量等于所有特征取值数量之和乘维度,在超大规模特征场景下内存经常爆掉。业界常见的做法是把所有特征的 Embedding 合在一张巨型表里,用一个特征索引偏移量去定位。这种做法在单机训练时能省不少内存,在分布式场景下更是标配。为了让代码简洁易懂,下面的实现还是按分表方式写的,生产环境落地时建议合并成大表。
3.3 FM 一阶部分与二阶交叉部分的实现
一阶部分有两种做法。如果特征全是类别型,可以直接用nn.Embedding(num_features, 1)来学每个特征的偏置,再对每个样本把它的特征偏置求和。但实际场景里通常还有数值型特征,这时最简单的方式是:
- 类别特征的一阶权重用
nn.Embedding - 数值特征直接套一个线性层,输入维度等于数值特征个数
二阶部分按我们前面说的“平方和减和平方”的化简式来实现:
class FM(nn.Module): def __init__(self, feat_num, embed_dim): super().__init__() self.feat_num = feat_num self.embed_dim = embed_dim def forward(self, embeddings): # embeddings shape: (batch, feat_num, embed_dim) sum_squared = torch.sum(embeddings, dim=1) ** 2 squared_sum = torch.sum(embeddings ** 2, dim=1) second_order = 0.5 * (sum_squared - squared_sum).sum(dim=1, keepdim=True) return second_order很多人在实现这一步时会踩一个坑:忘记对dim=1(特征维度)而不是dim=0(batch 维度)做求和,导致输出的二阶项维度错误。而且求和出来的是 (batch, embed_dim) 形状,再对最后一维求和,才能得到每个样本的一个标量。
3.4 Deep 部分的网络设计
Deep 部分的输入是“所有稀疏特征的 Embedding 拼接 + 数值特征拼接”。所以训练时,先对每个稀疏特征查表,把得到的 (batch, embed_dim) 表示在特征维度上拼起来,形成 (batch, feat_num * embed_dim) 的稠密向量,再和数值特征拼接,输入到全连接网络。
class DeepPart(nn.Module): def __init__(self, input_dim, hidden_dims, dropout=0.2): super().__init__() layers = [] in_dim = input_dim for h in hidden_dims: layers.append(nn.Linear(in_dim, h)) layers.append(nn.BatchNorm1d(h)) layers.append(nn.ReLU()) layers.append(nn.Dropout(dropout)) in_dim = h self.mlp = nn.Sequential(*layers) def forward(self, x): return self.mlp(x)BatchNorm 在 CTR 模型里的作用经常被低估。特征尺度差异大的时候,BatchNorm 不仅能加速收敛,还能缓解一些梯度爆炸的问题。但要注意,BN 的统计量依赖 batch 大小,如果线上推理时 batch 太小,效果会有波动,需要冻结 BN 的 running stats。
3.5 组装完整 DeepFM 模型
把前面的 FM 部分和 Deep 部分合并,最后把一阶输出、二阶输出、Deep 输出相加,经过 Sigmoid 得到点击概率:
class DeepFM(nn.Module): def __init__(self, feat_num, feat_size_list, embed_dim, hidden_dims, dense_dim): super().__init__() self.feat_num = feat_num self.embedding_layers = nn.ModuleList([ nn.Embedding(feat_size, embed_dim) for feat_size in feat_size_list ]) # 一阶:类别特征的偏置 + 数值特征的线性层 self.fc_first_sparse = nn.ModuleList([ nn.Embedding(feat_size, 1) for feat_size in feat_size_list ]) self.fc_first_dense = nn.Linear(dense_dim, 1) self.fm = FM(feat_num, embed_dim) deep_input_dim = feat_num * embed_dim + dense_dim self.deep = DeepPart(deep_input_dim, hidden_dims) self.deep_out = nn.Linear(hidden_dims[-1], 1) def forward(self, dense_x, sparse_x): # sparse_x shape: (batch, feat_num) embeddings = [emb(sparse_x[:, i]) for i, emb in enumerate(self.embedding_layers)] embeddings = torch.stack(embeddings, dim=1) # (batch, feat_num, embed_dim) first_sparse = sum([emb(sparse_x[:, i]) for i, emb in enumerate(self.fc_first_sparse)]) first_dense = self.fc_first_dense(dense_x) first_order = first_sparse + first_dense second_order = self.fm(embeddings) deep_in = torch.cat([embeddings.view(embeddings.size(0), -1), dense_x], dim=1) deep_out = self.deep_out(self.deep(deep_in)) y = first_order + second_order + deep_out return torch.sigmoid(y)3.6 损失函数和评估指标的选择
CTR 预估是二分类问题,损失函数用 BCE(二元交叉熵)就对了:
criterion = nn.BCELoss() optimizer = torch.optim.Adam(model.parameters(), lr=0.001)优化器这块,我通常先用 Adam 跑,收敛速度快,调参成本低。等模型稳定以后,如果想再提升一点精度,可以换 SGDM 配合学习率衰减,但需要更多时间调参,收益也比较有限。评估指标除了传统的 LogLoss(BCE loss 本身),强烈建议同时看 AUC。AUC 度量的是模型的排序能力,不依赖于具体阈值,在正负样本不平衡的情况下比准确率可靠得多。PyTorch 没有内置 AUC 计算,可以借助 sklearn 的roc_auc_score。
4. 训练完整流程与参数调优实测
4.1 训练主流程的完整代码
训练主流程没什么特别的,就是常规的 epoch 循环加 batch 迭代。但有一个细节值得提醒:nn.Embedding默认会参与梯度更新,如果某个特征取值在某个 batch 里没有出现,对应那一行的梯度就是 0,虽然不会报错,但更新效率会偏低。所以类别特征的 batch shuffle 要做得充分,让每个特征取值在训练过程中有均匀的曝光机会。
model.train() for epoch in range(num_epochs): for batch_dense, batch_sparse, batch_label in train_loader: pred = model(batch_dense, batch_sparse) loss = criterion(pred, batch_label) optimizer.zero_grad() loss.backward() optimizer.step()每跑完一个 epoch,在验证集上算一下 AUC 和 LogLoss,记录下来对比,判断是继续训练还是提前停止。这里我常用的判断标准是:连续两个 epoch 验证集 AUC 没有提升,就提前终止训练,同时把之前最优的模型参数保存下来。
4.2 几组关键超参的调优经验
Embedding 维度:从 4 开始试,每次翻倍,对比 AUC。在大多数中小规模数据集上,8 到 16 之间就能看到明显的收益拐点,继续加大收益甚微、内存翻倍。
Deep 隐藏层结构:试过 [64]、[128, 64]、[256, 128, 64] 等配置,实际数据显示从单层升到两层时,AUC 提升约 0.3% 到 0.5%,从两层到三层提升不足 0.1%。所以如果你数据集不够大,别一味加深网络。
Dropout:建议设置 0.2 到 0.5 之间。CTR 特征噪声大,过拟合现象非常普遍。我实测过一组对比,训练集 AUC 0.91、验证集 0.84 变成训练集 0.88、验证集 0.86,这就是加了 Dropout 之后的效果。
学习率:Adam 下 0.001 是通用起点,但如果用更大的 batch size(比如 4096),可以考虑开到 0.003。
4.3 训练过程中的收敛观察
以 Criteo 数据集的一个抽样子集为例,我训练了 10 个 epoch,batch size 128。前 3 个 epoch 里,训练 loss 从 0.693 快速回落到 0.52 左右,验证集 AUC 从 0.71 升到 0.78。第 4 到第 7 个 epoch,loss 下降速度放缓,AUC 爬升到 0.80。第 7 个 epoch 之后,训练 loss 仍在降,但验证集 AUC 基本停滞,甚至有轻微下降,这就是过拟合信号。所以提前停止的 patience 设成 2 是合理的。
另外提一下,CTR 数据集经常是海量级别,几千万样本很常见。这时候别傻乎乎的每个 epoch 都全量过一遍,在线学习或者只抽近一周滑动窗口数据来训练的玩法也很常见。
5. 常见问题与避坑指南
5.1 训练 loss 不下降是怎么回事
最可能的原因是类别特征索引没有做连续编码。nn.Embedding只接受值域在[0, num_embeddings-1]范围内的整数索引,如果特征编码是从 1 开始的、中间还有断层,模型要么报错要么学的 Embedding 是碎片化的。在编码阶段就检查:
sparse_feat = sparse_feat.astype('category').cat.codes确保没有 -1 或者超过num_embeddings的值。写个断言过滤一下。
5.2 Embedding 维度与特征数量不匹配
训练时如果发现显存占用涨得夸张,先检查是不是某个稀疏特征的取值数量太大,比如用户 ID 有上亿个不同的值,那 Embedding 表的参数量就是 1 亿乘以维度,在一张普通显卡上直接 OOM。解决办法是:对高频特征保留完整 Embedding,低频长尾特征全部映射到一个共享的“未知”向量里。业界叫 hash trick 或者 frequency filtering。这个操作在工程上几乎必做。
5.3 模型过拟合怎么判断和处理
CTR 场景特征极度稀疏,模型很容易记住训练集里的噪声,过拟合是常态,不用慌。处理手段除了前面提到的 Dropout,还有一个很有效的做法是给 Embedding 层加 L2 正则:
optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-5)weight_decay加在 Embedding 上效果立竿见影,但注意不要加太大,否则特征向量被压扁,表达力下降,AUC 反而会掉。
5.4 离线指标好但线上效果差
这个问题的核心原因通常是训练和推理时特征分布不一致,俗称线上线下不一致。常见的坑是线上特征实时拼接时某个字段缺失,填了 0 或者默认值,但离线训练时这个特征是有真实值的。要尽量模拟线上特征处理流程来做离线训练数据,比如做成一个统一的特征工程 pipeline,训练和推理共用同一套逻辑。
另一个容易忽略的点是样本偏差。如果用“曝光 + 点击”来构造样本,没曝光的优质物品就没有机会进入训练集,模型会有选择偏差。这个问题超出了 DeepFM 本身的范畴,但做推荐系统迟早会遇到,提前知道能省很多事。
5.5 训练速度太慢怎么办
数据量大时 PyTorch 的 DataLoader 可能成为瓶颈。加num_workers=4或8以外,还可以把类别特征和数值特征打包成 TensorDataset,减少 Python 层面的额外开销。如果模型本身重,Deep 部分的算力耗费比 FM 部分大得多,可以考虑减少隐藏层宽度。
还有一个优化点:nn.Embedding(sparse=True)在 CPU 上对大规模特征是友好的,但在 GPU 上和普通 Embedding 的性能差别不大,如果不需要稀疏梯度,直接用常规 Embedding 即可。
6. 从 Demo 到线上部署还要做什么
在本地把 DeepFM 跑通只是第一步,如果要上生产环境,后续还有几个关键环节要处理。
模型导出方面,PyTorch 的模型可以用 TorchScript 导出,这样线上推理不依赖 Python 的完整环境,用 C++ 也能加载,时延有明显下降。导出的核心是把模型结构固化成计算图:
traced_model = torch.jit.script(model) traced_model.save("deepfm.pt")另外,CTR 模型线上服务通常有严格的性能要求,单个请求的推理时延不能太高。对 DeepFM 来说,最大的性能瓶颈在特征处理和 Embedding 查询。业界常用做法是把训练好的 Embedding 向量缓存到 Redis 之类的线上存储里,在线推理时直接查询拿到稠密向量,再进 Deep 网络做前向计算。
如果你的特征规模没有大到必须做特征服务拆分,也可以直接复用训练时的模型加载方式,但务必压测确认 QPS 和时延符合需求。
在实际动手复现 DeepFM 的过程中,我最大的体会是:模型结构本身并不复杂,真正的难点在于处理特征的细节、超参的配合,以及从离线到线上的一致性保障。很多人看了论文觉得懂了,但一上手跑数据就发现各种问题,原因就在于这些工程细节。建议你拿到代码之后,先在小数据集上完整跑通一遍训练和评估流程,再逐步放大特征量和数据量,遇到问题的时候回头对照这篇文章里的排查清单,大概率能省掉不少弯路。
本文还有配套的精品资源,点击获取