You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
121 lines
4.3 KiB
121 lines
4.3 KiB
4 years ago
|
import sys
|
||
|
|
import os.path as osp
|
||
|
|
this_dir = osp.dirname(__file__)
|
||
|
|
lib_path = osp.join(this_dir, '..')
|
||
|
|
sys.path.insert(0, lib_path)
|
||
|
|
from NCF.dataprocess import DataProcess
|
||
|
|
from NCF.network import GMF,MLP,NeuMF
|
||
|
|
from NCF.trainer import Trainer
|
||
|
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
|
||
|
|
gmf_config = {'num_epoch': 100,
|
||
|
|
'batch_size': 1024,
|
||
|
|
'optimizer': 'adam',
|
||
|
|
'adam_lr': 1e-3,
|
||
|
|
'num_users': 6040,
|
||
|
|
'num_items': 3706,
|
||
|
|
'latent_dim_gmf': 8,
|
||
|
|
'num_negative': 4,
|
||
|
|
'layers': [],
|
||
|
|
'l2_regularization': 0, # 0.01
|
||
|
|
'pretrain': False, # do not modify this
|
||
|
|
'use_cuda': True,
|
||
|
|
'device_id': 0,
|
||
|
|
'model_name': '../TrainedModels/NCF_GMF.model'
|
||
|
|
}
|
||
|
|
|
||
|
|
mlp_config = {'num_epoch': 100,
|
||
|
|
'batch_size': 1024, # 1024,
|
||
|
|
'optimizer': 'adam',
|
||
|
|
'adam_lr': 1e-3,
|
||
|
|
'num_users': 6040,
|
||
|
|
'num_items': 3706,
|
||
|
|
'latent_dim_mlp': 8,
|
||
|
|
'latent_dim_gmf': 8,
|
||
|
|
'num_negative': 4,
|
||
|
|
'layers': [16,64,32,16,8], # layers[0] is the concat of latent user vector & latent item vector
|
||
|
|
'l2_regularization': 0.0000001, # MLP model is sensitive to hyper params
|
||
|
|
'use_cuda': True,
|
||
|
|
'device_id': 0,
|
||
|
|
'pretrain': True,
|
||
|
|
'gmf_config': gmf_config,
|
||
|
|
'pretrain_gmf': '../TrainedModels/NCF_GMF.model',
|
||
|
|
'model_name': '../TrainedModels/NCF_MLP.model'
|
||
|
|
}
|
||
|
|
|
||
|
|
neumf_config = {'num_epoch': 100,
|
||
|
|
'batch_size': 1024, #1024
|
||
|
|
'optimizer': 'adam',
|
||
|
|
'adam_lr': 1e-3,
|
||
|
|
'num_users': 6040,
|
||
|
|
'num_items': 3706,
|
||
|
|
'latent_dim_gmf': 8,
|
||
|
|
'latent_dim_mlp': 8,
|
||
|
|
'num_negative': 4,
|
||
|
|
'layers': [16,32,16,8], # layers[0] 是用户和物品隐层表示concat的维度
|
||
|
|
'l2_regularization': 0.01,
|
||
|
|
'alpha': 0.5, # 用于控制GMF和MLP模型参数的权重
|
||
|
|
'use_cuda': True,
|
||
|
|
'device_id': 0,
|
||
|
|
'pretrain': False,
|
||
|
|
'gmf_config': gmf_config,
|
||
|
|
'pretrain_gmf': '../TrainedModels/NCF_GMF.model',
|
||
|
|
'mlp_config': mlp_config,
|
||
|
|
'pretrain_mlp': '../TrainedModels/NCF_MLP.model',
|
||
|
|
'model_name': '../TrainedModels/NCF_NeuMF.model'
|
||
|
|
}
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
####################################################################################
|
||
|
|
# NCF 神经协同过滤算法
|
||
|
|
####################################################################################
|
||
|
|
|
||
|
|
# 加载和预处理数据
|
||
|
|
dp = DataProcess("../Data/ml-1m/ratings.dat")
|
||
|
|
|
||
|
|
# 初始化GMP模型
|
||
|
|
config = gmf_config
|
||
|
|
model = GMF(config, config['latent_dim_gmf'])
|
||
|
|
|
||
|
|
# # 初始化MLP模型
|
||
|
|
config = mlp_config
|
||
|
|
model = MLP(config, config['latent_dim_mlp'])
|
||
|
|
|
||
|
|
# 初始化NeuMF模型
|
||
|
|
config = neumf_config
|
||
|
|
model = NeuMF(config, config['latent_dim_gmf'], config['latent_dim_mlp'])
|
||
|
|
|
||
|
|
# ###############################################################
|
||
|
|
# 模型训练阶段
|
||
|
|
# ###############################################################
|
||
|
|
trainer = Trainer(model=model, config=config)
|
||
|
|
trainer.train(dp.sample_generator)
|
||
|
|
trainer.save()
|
||
|
|
|
||
|
|
# ###############################################################
|
||
|
|
# 模型测试阶段
|
||
|
|
# ###############################################################
|
||
|
|
|
||
|
|
# 加载数据集
|
||
|
|
# dp = DataProcess("../Data/ml-1m/ratings.dat")
|
||
|
|
|
||
|
|
config = neumf_config
|
||
|
|
neumf = NeuMF(config, config['latent_dim_gmf'], config['latent_dim_mlp'])
|
||
|
|
state_dict = torch.load("../TrainedModels/NCF_NeuMF.model", map_location=torch.device('cpu'))
|
||
|
|
neumf.load_state_dict(state_dict, strict=False)
|
||
|
|
|
||
|
|
# 对用户User_id喜好度进行预测
|
||
|
|
User_id = 1
|
||
|
|
result = np.zeros((3706))
|
||
|
|
for j in range(3706):
|
||
|
|
socre = neumf.forward(torch.LongTensor([User_id]), torch.LongTensor([j]))
|
||
|
|
score = socre.detach().numpy()
|
||
|
|
result[j] = socre[0][0]
|
||
|
|
# 选取User_id喜好度最高的N个电影id进行推荐
|
||
|
|
N = 5
|
||
|
|
indexs = np.argsort(-result)[:N]
|
||
|
|
print(indexs)
|
||
|
|
|
||
|
|
|