基于 Criteo 1M 数据集的 CTR 预估 -- 模型训练部分

前言

在前面一篇文章完成了 Criteo 数据集的清洗与特征处理，得到了DataLoader 类型的数据批，接下来就可以选择合适的神经网络进行 CTR 预估了，博主打算检验一下自己学的经典推荐模型，因此后面会使用多个模型进行训练，同时也横向对各个模型做个对比，这次先使用 DeepFM模型

DeepFM 模型

FM部分：
FM 的公式如下：

公式的前两项是偏置项 w₀以及一阶项，FM 的精髓在于二阶交叉特征项，也就是最后一项，x_i与 x_j表示特征 i 与特征 j 的值，前面的 <v_i, v_j> 是这两个特征对应的权重，显而易见的是把 v_i与 v_j存储在一个 n x n 的矩阵 W 中（n 为特征数量），但特征数量很多时会导致矩阵占用空间巨大，FM 的做法是利用 W = V^TV，V的形状为 n x k，k 是超参数，可人为选择，这样模型仅需存储一个矩阵 V，从而降低了模型的空间复杂度
在我的模型中，没有使用偏置以及一阶项，因为二阶特征交叉已经包括了足够多的信息，从深度学习的角度理解，V 矩阵实际就是 Embedding 层的权重矩阵
在实际计算中，二阶特征交互项通常使用下列等价公式来优化时间复杂度：

具体推导大家如果有兴趣可以去看看，我就不写出来了
Deep 部分：
Deep 部分从上面的结构图中看就可以看出来是一个 MLP，即多层感知机，由多个全连接层（图中画了 3 层，仅为示例）和激活函数组成，大家应该都十分熟悉，也就不多介绍了

模型实现

classDeepFM(nn.Module):def__init__(self,emb_size_of_every_col,feature_cols,dim):super(DeepFM,self).__init__()# 每一列一个 Embedding 列的先后顺序与 all_cat_cols 中一致self.embeddings=nn.ModuleList([nn.Embedding(size,dim,max_norm=1)forsizeinemb_size_of_every_col])self.mlp=self._mlp(feature_cols*dim)def_mlp(self,dim):returnnn.Sequential(nn.Linear(dim,dim//2),nn.Tanh(),nn.Linear(dim//2,dim//4),nn.Tanh(),nn.Linear(dim//4,1)# 不需要 sigmoid)defDeep(self,features_embs):# features_embs: [batch_size, feature_cols, dim]features_embs=features_embs.reshape((features_embs.shape[0],-1))# [batch_size, feature_cols * dim]returnself.mlp(features_embs)defFMCross(self,features_embs):# features_embs: [batch_size, feature_cols, dim]square_of_sum=torch.sum(features_embs,dim=1)**2# [batch_size, dim]sum_of_square=torch.sum(features_embs**2,dim=1)# [batch_size, dim]return0.5*torch.sum(square_of_sum-sum_of_square,dim=1)# [batch_size]defforward(self,X):# X: [batch_size, feature_cols]num_cols=X.shape[1]features=[]# 取 X 的每一列特征，送入该列的 Embeddingforiinrange(num_cols):feature=self.embeddings[i](X[:,i])# [batch_size, 1, dim] 第 i 列特征的 Embedding 向量features.append(feature)features=torch.stack(features,dim=1)# [batch_size, feature_cols, dim]# Deep 部分uv_deep=self.Deep(features)# [batch_size, 1]uv_deep=torch.squeeze(uv_deep,dim=1)# [batch_size]# FM 部分uv_fm=self.FMCross(features)# [batch_size]returnuv_deep+uv_fm# 由于损失函数采用 BCEWithLogitsLoss 所以不需要取 sigmoid

上面是 DeepFM 模型的 PyTorch 实现，前面的特征处理中，我们得到了39列类别特征，包括原始的26列类别特征以及对整数特征分桶得到的13列类别特征，这些列的含义各不相同，因此模型内维护了一个 nn.ModuleList，为每一个特征列都提供 Embedding，模型初始化时要传入每列的词表大小 emb_size_of_every_col 用于初始化各列的 Embedding
其余代码则按照模型结构图实现

训练循环

模型编写好就可以来写训练循环了，首先加载数据，这些在上一篇博客里已经写了：

train_df,valid_df,test_df,all_cat_cols,encoders=process(data_path)train_iter,valid_iter,test_iter=load_data(train_df,valid_df,test_df,all_cat_cols,batch_size,num_workers)emb_size_of_every_col=[len(encoders[col])+1forcolinall_cat_cols]# 每个特征列特征值的最大值（这一列 Embedding 的大小）

接着定义模型、损失函数以及优化器

# 网络定义model=DeepFM(emb_size_of_every_col,len(all_cat_cols),dim=dim)criterion=nn.BCEWithLogitsLoss()optimizer=torch.optim.AdamW(model.parameters(),lr=lr,weight_decay=weight_decay)

损失函数采用二元交叉熵损失函数 BCEWithLogitsLoss() 该函数在计算损失时会自动对输入的张量求 sigmoid，这也是模型 MLP 的输出以及总输出没有取 sigmoid 的原因
优化器采用 AdamW，包含权重衰减参数，能缓解过拟合
接下来开始编写训练循环，训练基本策略如下：
训练集：每个 epoch 遍历 train_iter，进行前向传播、算 loss、反向传播、优化器更新；计算模型在训练集上的损失
验证集：每个 epoch 训练完后跑一次验证集 valid_iter，注意使用上下文管理器 torch.no_grad() 不更新参数；计算验证集的损失、AUC；根据验证集 AUC 保存最佳模型
测试集：训练全部结束后，加载验证集上最好模型，仅跑一次 test_iter，得到最终 AUC
模型的训练以及参数更新：

foreinrange(epochs):# 模型训练model.train()total_train_loss=0.0train_samples=0forX,yintrain_iter:optimizer.zero_grad()loss=criterion(model(X),y)loss.backward()optimizer.step()total_train_loss+=loss.item()*len(y)train_samples+=len(y)train_loss=total_train_loss/train_samples

每一个 epoch 结束后，在验证集上计算 avg_loss 以及 AUC

valid_loss,auc=evaluate(model,valid_iter,criterion)

evaluate 函数的实现：

defevaluate(model,valid_iter,criterion):model.eval()all_y_pred=[]all_y_label=[]total_loss=0.0num_samples=0withtorch.no_grad():forX,y_labelinvalid_iter:logits=model(X)loss=criterion(logits,y_label)y_pred=torch.sigmoid(logits)# 模型输出的是 logits，未经过 sigmoid# 计算总损失与样本数，后面会计算 avg_losstotal_loss+=loss.item()*len(y_label)# BCEWithLogitsLoss 返回的是误差均值，别忘了乘该批次样本数量num_samples+=len(y_label)# 将 y_pred 与 y_label 转为 numpy 数组并加入列表，后面方便计算 aucall_y_pred.append(y_pred.cpu().numpy())all_y_label.append(y_label.cpu().numpy())# 计算损失avg_loss=total_loss/num_samples# 计算 AUCy_pred=np.concatenate(all_y_pred).reshape(-1)y_label=np.concatenate(all_y_label).reshape(-1)auc=roc_auc_score(y_label,y_pred)returnavg_loss,auc

计算 AUC 使用的是 sklearn 的 roc_auc_score 函数，传入预测值和真实值的 numpy 数组即可
根据计算的 auc 更新最佳模型以及 best_auc，同时加入早停策略

ifauc>best_auc:# 更新最佳模型best_auc=auc torch.save(model.state_dict(),"save/best_deepfm.pt")bad_epochs=0else:# 如果 AUC 没有上涨，则早停计数加一bad_epochs+=1ifbad_epochs>=early_stop_patience:print(f"Training early stopped at epoch{e+1}")break

训练所有 epoch 结束后，加载最佳模型，在测试集上计算 AUC

model.load_state_dict(torch.load("save/best_deepfm.pt"))_,auc=evaluate(model,test_iter,criterion)print(f"Test AUC:{auc:.4f}")

完整的训练代码如下：

importnumpyasnpimporttorchimporttorch.nnasnnfromsklearn.metricsimportroc_auc_scorefromdata_process.processimportprocessfromdata_process.loadimportload_datafrommodels.DeepFMimportDeepFM# 采用 Kaggle 上的 Criteo_1M_with_nans.csv 数据集data_path="data/Criteo_1M_with_nans.csv"batch_size=256num_workers=4# 使用子进程数dim=16# Embedding 层的维度lr=1e-3weight_decay=1e-4defevaluate(model,valid_iter,criterion):model.eval()all_y_pred=[]all_y_label=[]total_loss=0.0num_samples=0withtorch.no_grad():forX,y_labelinvalid_iter:logits=model(X)loss=criterion(logits,y_label)y_pred=torch.sigmoid(logits)# 模型输出的是 logits，未经过 sigmoid# 计算总损失与样本数，后面会计算 avg_losstotal_loss+=loss.item()*len(y_label)# BCEWithLogitsLoss 返回的是误差均值，别忘了乘该批次样本数量num_samples+=len(y_label)# 将 y_pred 与 y_label 转为 numpy 数组并加入列表，后面方便计算 aucall_y_pred.append(y_pred.cpu().numpy())all_y_label.append(y_label.cpu().numpy())# 计算损失avg_loss=total_loss/num_samples# 计算 AUCy_pred=np.concatenate(all_y_pred).reshape(-1)y_label=np.concatenate(all_y_label).reshape(-1)auc=roc_auc_score(y_label,y_pred)returnavg_loss,aucdefmain(epochs,early_stop_patience,best_auc,bad_epochs):train_df,valid_df,test_df,all_cat_cols,encoders=process(data_path)train_iter,valid_iter,test_iter=load_data(train_df,valid_df,test_df,all_cat_cols,batch_size,num_workers)emb_size_of_every_col=[len(encoders[col])+1forcolinall_cat_cols]# 每个特征列特征值的最大值（这一列 Embedding 的大小）# 网络定义model=DeepFM(emb_size_of_every_col,len(all_cat_cols),dim=dim)criterion=nn.BCEWithLogitsLoss()optimizer=torch.optim.AdamW(model.parameters(),lr=lr,weight_decay=weight_decay)foreinrange(epochs):# 模型训练model.train()total_train_loss=0.0train_samples=0forX,yintrain_iter:optimizer.zero_grad()loss=criterion(model(X),y)loss.backward()optimizer.step()total_train_loss+=loss.item()*len(y)train_samples+=len(y)train_loss=total_train_loss/train_samples# 每一个 epoch 结束后，在验证集上计算 avg_loss 以及 AUCvalid_loss,auc=evaluate(model,valid_iter,criterion)print(f"Epoch{e+1}: \ntrain loss:{train_loss:.4f}\nvalid loss:{valid_loss:.4f}\nvalid AUC:{auc:.4f}")ifauc>best_auc:# 更新最佳模型best_auc=auc torch.save(model.state_dict(),"save/best_deepfm.pt")bad_epochs=0else:# 如果 AUC 没有上涨，则早停计数加一bad_epochs+=1ifbad_epochs>=early_stop_patience:print(f"Training early stopped at epoch{e+1}")break# 训练结束，加载最佳模型，在测试集上计算 AUCmodel.load_state_dict(torch.load("save/best_deepfm.pt"))_,auc=evaluate(model,test_iter,criterion)print(f"Test AUC:{auc:.4f}")if__name__=="__main__":main(epochs=10,early_stop_patience=3,best_auc=0.0,bad_epochs=0)

模型参数以及训练参数的设定（非最优，后面可能还会继续调整）：