DeepCTR-Torch 如何保存模型

最新推荐文章于 2025-11-13 11:01:55 发布

原创

最新推荐文章于 2025-11-13 11:01:55 发布 · 2.7k 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#python

这篇博客介绍了在使用DeepCTR库时遇到的模型保存和加载问题。当尝试用TensorFlow的`save_model`和`load_model`函数保存和加载DeepFM模型时，出现AttributeError。解决方法是改用PyTorch的方式，通过`torch.save`和`torch.load`来保存和加载模型的状态字典。示例代码展示了如何在DeepFM模型上实现这一过程。

官网给出的示例：

DeepCTR Documentation, Release 0.9.0
from tensorflow.python.keras.models import save_model,load_model
model = DeepFM()
save_model(model, 'DeepFM.h5')# save_model, same as before
from deepctr.layers import custom_objects
model = load_model('DeepFM.h5',custom_objects)# load_model,just add a parameter

实际执行的时候会报错：

AttributeError: 'DeepFM' object has no attribute 'outputs'

解决办法：

采用 torch的模型保存办法，下述示例是 Deep-torch 的示例文件 run_classification_criteo.py

from deepctr_torch.models import *

#创建model1

model = DeepFM(linear_feature_columns, dnn_feature_columns, task='binary')

#model 处理

#...

保存以及读取模型：

#model 保存

import torch

torch.save(model.state_dict(),"a.txt")

model2 = DeepFM(linear_feature_columns, dnn_feature_columns, task='binary')

model2.load_state_dict(torch.load("a.txt"))

pred_ans2 = model2.predict(test_model_input, batch_size=25