import torch
import torch.nn as nn
import torch.sparse as sparse
import torch.nn.functional as F
import torch.optim as optim
import csv
import os
import sys
import datetime
import math
from Models import *
from main import load_pretrained_data
from utility.parser import parse_args
from utility.helper import *
from utility.batch_test import *
class Model_Wrapper(object):
def __init__(self, data_config, pretrain_data):
"""
初始化 Model_Wrapper 类的实例
参数:
- data_config (dict): 包含数据相关配置的字典,例如用户数量、物品数量和规范化邻接矩阵等
- pretrain_data: 预训练数据(如果有)
以下是对各个属性的初始化和设置:
"""
# 从命令行参数获取模型类型
self.model_type = args.model_type
# 从命令行参数获取邻接矩阵类型
self.adj_type = args.adj_type
# 从命令行参数获取算法类型
self.alg_type = args.alg_type
# 评估消息丢弃的设置
self.mess_dropout = eval(args.mess_dropout)
#eval把字符串转变为python表达式进行求解
# 保存预训练数据
self.pretrain_data = pretrain_data
# 获取数据配置中的用户数量
self.n_users = data_config['n_users']
# 获取数据配置中的物品数量
self.n_items = data_config['n_items']
# 获取规范化的邻接矩阵,并将其转换为 PyTorch 的稀疏张量并转换为浮点数类型
self.norm_adj = data_config['norm_adj']
self.norm_adj = self.sparse_mx_to_torch_sparse_tensor(self.norm_adj).float()
# 记录阿尔法值的标志,默认为 False
self.record_alphas = False
# 学习率设置
self.lr = args.lr
# 嵌入维度设置
self.emb_dim = args.embed_size
# 批量大小设置
self.batch_size = args.batch_size
# 评估层大小并将其转换为列表
self.weight_size = eval(args.layer_size)
# 获取层数
self.n_layers = len(self.weight_size)
# 构建模型类型的字符串表示,包含邻接矩阵类型、算法类型和层数等信息
self.model_type += '_%s_%s_l%d' % (self.adj_type, self.alg_type, self.n_layers)
# 评估正则化系数并将其转换为列表
self.regs = eval(args.regs)
# 获取第一个正则化系数作为衰减系数
self.decay = self.regs[0]
# 详细输出的标志
self.verbose = args.verbose
# 打印模型的类型
print('model_type is {}'.format(self.model_type))
# 构建模型权重的保存路径
self.weights_save_path = '%sweights/%s/%s/l%s/r%s' % (args.weights_path, args.dataset, self.model_type,
str(args.lr), '-'.join([str(r) for r in eval(args.regs)]))
"""
打印当前算法类型,并根据算法类型创建相应的模型
"""
print('----self.alg_type is {}----'.format(self.alg_type))
# 如果算法类型是 'ngcf',创建 NGCF 模型
if self.alg_type in ['ngcf']:
self.model = NGCF(self.n_users, self.n_items, self.emb_dim, self.weight_size, self.mess_dropout)
# 如果算法类型是'mf',创建 MF 模型
elif self.alg_type in ['mf']:
self.model = MF(self.n_users, self.n_items, self.emb_dim)
# 如果算法类型不在上述列表中,抛出异常
else:
raise Exception('Dont know which model to train')
# 将模型移动到 GPU 进行计算
self.model = self.model.cuda()
# 将规范化的邻接矩阵也移动到 GPU
self.norm_adj = self.norm_adj.cuda()
# 创建优化器,使用 Adam 算法优化模型的参数,学习率为设置的 self.lr
self.optimizer = optim.Adam(self.model.parameters(), lr=self.lr)
# 设置学习率调度器
self.lr_scheduler = self.set_lr_scheduler()
# 打印模型的结构
print(self.model)
# 打印模型中每个参数的名称和大小
for name, param in self.model.named_parameters():
print(name,' ', param.size())
class MyModel:
"""
自定义的模型类,包含了设置学习率调度器、保存和加载模型、测试和训练模型的方法
"""
def set_lr_scheduler(self):
"""
设置学习率调度器
定义了一个函数 `fac`,它根据传入的 epoch 计算学习率的缩放因子
使用 `optim.lr_scheduler.LambdaLR` 创建学习率调度器,并返回
Returns:
torch.optim.lr_scheduler._LRScheduler: 学习率调度器对象
"""
# 定义一个匿名函数 fac,根据 epoch 计算学习率缩放因子
fac = lambda epoch: 0.96 ** (epoch / 50)
# 创建 LambdaLR 学习率调度器,根据 fac 函数调整学习率
scheduler = torch.optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda=fac)
return scheduler
def save_model(self):
"""
保存模型
确保保存模型的路径存在,然后使用 `torch.save` 保存模型的状态字典
"""
# 确保保存模型的路径存在
ensureDir(self.weights_save_path)
# 保存模型的状态字典到指定路径
torch.save(self.model.state_dict(), self.weights_save_path)
def load_model(self):
"""
加载模型
从指定路径加载模型的状态字典
"""
# 从指定路径加载模型的状态字典
self.model.load_state_dict(torch.load(self.weights_save_path))
def test(self, users_to_test, drop_flag=False, batch_test_flag=False):
"""
测试模型
将模型设置为评估模式,计算用户和物品的嵌入,然后进行测试并返回结果
Args:
users_to_test (list): 要测试的用户列表
drop_flag (bool, optional): 标志,默认为 False
batch_test_flag (bool, optional): 标志,默认为 False
Returns:
dict: 测试结果
"""
# 将模型设置为评估模式
self.model.eval()
with torch.no_grad():
# 通过模型计算用户和物品的嵌入
ua_embeddings, ia_embeddings = self.model(self.norm_adj)
# 进行测试并返回结果
result = test_torch(ua_embeddings, ia_embeddings, users_to_test)
return result
def train(self):
"""
训练模型
包含了训练的整个流程,包括数据采样、前向传播、计算损失、反向传播、学习率调整、评估和保存模型等
"""
# 初始化训练时间列表
training_time_list = []
# 初始化各种指标的记录列表
loss_loger, pre_loger, rec_loger, ndcg_loger, hit_loger, map_loger, mrr_loger, fone_loger = [], [], [], [], [], [], [], []
# 初始化早停相关的变量
stopping_step = 0
should_stop = False
cur_best_pre_0 = 0.
# 计算训练的批次数
n_batch = data_generator.n_train // args.batch_size + 1
for epoch in range(args.epoch):
# 记录训练开始时间
t1 = time.time()
# 初始化每个 epoch 的损失
loss, mf_loss, emb_loss, reg_loss = 0., 0., 0., 0.
n_batch = data_generator.n_train // args.batch_size + 1
f_time, b_time, loss_time, opt_time, clip_time, emb_time = 0., 0., 0., 0., 0., 0.
sample_time = 0.
cuda_time = 0.
for idx in range(n_batch):
# 将模型设置为训练模式
self.model.train()
# 优化器梯度清零
self.optimizer.zero_grad()
# 记录采样开始时间
sample_t1 = time.time()
# 从数据生成器采样用户、正样本物品和负样本物品
users, pos_items, neg_items = data_generator.sample()
# 计算采样时间
sample_time += time.time() - sample_t1
# 通过模型计算用户和物品的嵌入
ua_embeddings, ia_embeddings = self.model(self.norm_adj)
# 获取特定用户的嵌入
u_g_embeddings = ua_embeddings[users]
# 获取正样本物品的嵌入
pos_i_g_embeddings = ia_embeddings[pos_items]
# 获取负样本物品的嵌入
neg_i_g_embeddings = ia_embeddings[neg_items]
# 计算批量的损失
batch_mf_loss, batch_emb_loss, batch_reg_loss = self.bpr_loss(u_g_embeddings, pos_i_g_embeddings,
neg_i_g_embeddings)
# 计算总的批量损失
batch_loss = batch_mf_loss + batch_emb_loss + batch_reg_loss
# 反向传播
batch_loss.backward()
# 优化器更新参数
self.optimizer.step()
#self.optimizer.step()
# 累计损失
loss += float(batch_loss)
# print('loss: ', loss)
mf_loss += float(batch_mf_loss)
emb_loss += float(batch_emb_loss)
reg_loss += float(batch_reg_loss)
# 执行学习率调度器的一步更新
self.lr_scheduler.step()
# 释放一些变量以节省内存
del ua_embeddings, ia_embeddings, u_g_embeddings, neg_i_g_embeddings, pos_i_g_embeddings
# 检查损失是否为 NaN,如果是则报错并退出
if math.isnan(loss) == True:
print('ERROR: loss is nan.')
sys.exit()
# 每 10 个 epoch 打印测试评估指标
if (epoch + 1) % 10!= 0:
if args.verbose > 0 and epoch % args.verbose == 0:
# 构建性能字符串
perf_str = 'Epoch %d [%.1fs]: train==[%.5f=%.5f + %.5f]' % (
epoch, time.time() - t1, loss, mf_loss, emb_loss)
#training_time_list.append(time.time() - t1)
print(perf_str)
continue
# 记录测试开始时间
t2 = time.time()
# 获取要测试的用户列表
users_to_test = list(data_generator.test_set.keys())
# 进行测试并获取结果
ret = self.test(users_to_test, drop_flag=True)
# 记录训练时间
training_time_list.append(t2 - t1)
# 记录测试结束时间
t3 = time.time()
# 记录各种指标
loss_loger.append(loss)
rec_loger.append(ret['recall'])
pre_loger.append(ret['precision'])
ndcg_loger.append(ret['ndcg'])
hit_loger.append(ret['hit_ratio'])
map_loger.append(ret['map'])
mrr_loger.append(ret['mrr'])
fone_loger.append(ret['fone'])
if args.verbose > 0:
# 构建详细的性能字符串
perf_str = 'Epoch %d [%.1fs + %.1fs]: train==[%.5f=%.5f + %.5f + %.5f], recall=[%.5f, %.5f], '
'precision=[%.5f, %.5f], hit=[%.5f, %.5f], ndcg=[%.5f, %.5f], map=[%.5f, %.5f], mrr=[%.5f, %.5f], f1=[%.5f, %.5f]' % \
(epoch, t2 - t1, t3 - t2, loss, mf_loss, emb_loss, reg_loss, ret['recall'][0],
ret['recall'][-1],
ret['precision'][0], ret['precision'][-1], ret['hit_ratio'][0], ret['hit_ratio'][-1],
ret['ndcg'][0], ret['ndcg'][-1],ret['map'][0], ret['map'][-1],ret['mrr'][0], ret['mrr'][-1],ret['fone'][0], ret['fone'][-1])
print(perf_str)
# 进行早停判断和更新相关变量
cur_best_pre_0, stopping_step, should_stop = early_stopping(ret['recall'][0], cur_best_pre_0,
stopping_step, expected_order='acc',
flag_step=5)
# 早停条件:如果当前最佳精度连续下降 10 步
if should_stop:
break
# 如果当前召回率等于历史最佳召回率且保存标志为 1,保存模型
if ret['recall'][0] == cur_best_pre_0 and args.save_flag == 1:
# 保存模型
self.save_model()
if self.record_alphas:
self.best_alphas = [i for i in self.model.get_alphas()]
print('save the weights in path: ', self.weights_save_path)
# 如果需要保存推荐结果,将最终推荐结果打印到 CSV 文件
if args.save_recom:
results_save_path = '%soutput/%s/reRecommendation.csv' % (args.proj_path, args.dataset)
self.save_recResult(results_save_path)
# 如果召回率记录不为空,打印最终结果
if rec_loger!= []:
self.print_final_results(rec_loger, pre_loger, ndcg_loger, hit_loger, map_loger,mrr_loger,fone_loger,training_time_list)
def save_recResult(self, outputPath):
"""
此函数用于保存推荐结果
参数:
outputPath (str): 保存推荐结果的文件路径
"""
# 用于反转推荐列表
recommendResult = {}
u_batch_size = BATCH_SIZE * 2 # 用户的批处理大小
i_batch_size = BATCH_SIZE # 项目(物品)的批处理大小
# 获取所有要测试的用户
users_to_test = list(data_generator.test_set.keys())
n_test_users = len(users_to_test) # 测试用户的数量
n_user_batchs = n_test_users // u_batch_size + 1 # 用户批处理的数量
count = 0
# 计算结果
# 将模型设置为评估模式
self.model.eval()
with torch.no_grad():
# 获取用户和项目的嵌入
ua_embeddings, ia_embeddings = self.model(self.norm_adj)
# 按批处理获取结果
for u_batch_id in range(n_user_batchs):
# 计算批处理的起始和结束索引
start = u_batch_id * u_batch_size
end = (u_batch_id + 1) * u_batch_size
user_batch = users_to_test[start: end] # 获取当前批处理的用户
item_batch = range(ITEM_NUM) # 获取所有项目
# 获取用户和项目的嵌入
u_g_embeddings = ua_embeddings[user_batch]
i_g_embeddings = ia_embeddings[item_batch]
# 计算用户对项目的评分
#torch.matmul()表示的是pytorch中的矩阵乘法操作
rate_batch = torch.matmul(u_g_embeddings, torch.transpose(i_g_embeddings, 0, 1))
# 将数据从 GPU 移动到 CPU 并转换为 numpy 数组
rate_batch = rate_batch.detach().cpu().numpy()
# 将每个用户的评分与其用户 ID 关联
user_rating_uid = zip(rate_batch, user_batch)
# 为每个用户计算评分和推荐
for x in user_rating_uid:
# 用户 u 的评分
rating = x[0]
# 用户 ID
u = x[1]
training_items = data_generator.train_items[u] # 获取用户的训练项目
user_pos_test = data_generator.test_set[u] # 获取用户的测试集
all_items = set(range(ITEM_NUM)) # 所有项目的集合
test_items = list(all_items - set(training_items)) # 测试项目
item_score = {} # 用于存储每个测试项目的评分
for i in test_items:
item_score[i] = rating[i] # 存储项目的评分
K_max = max(Ks) # 获取最大的推荐数量
# 获取评分最高的 K_max 个项目
K_max_item_score = heapq.nlargest(K_max, item_score, key=item_score.get)
recommendResult[u] = K_max_item_score # 保存用户的推荐项目
# 确保输出路径存在
ensureDir(outputPath)
with open(outputPath, 'w') as f:
print("----the recommend result has %s items." % (len(recommendResult)))
for key in recommendResult.keys():
outString = ""
for v in recommendResult[key]:
outString = outString + "," + str(v)
f.write("%s%s\n"%(key,outString))
#f.write("%s,%s\n"%(key,recommendResult[key]))
class YourModel:
# 打印最终结果的方法
def print_final_results(self, rec_loger, pre_loger, ndcg_loger, hit_loger, map_loger, mrr_loger, fone_loger,
training_time_list):
"""
该方法用于打印和保存模型训练的最终结果
参数:
rec_loger: 召回率的日志列表
pre_loger: 准确率的日志列表
ndcg_loger: NDCG 的日志列表
hit_loger: 命中的日志列表
map_loger: MAP 的日志列表
mrr_loger: MRR 的日志列表
fone_loger: F1 值的日志列表
training_time_list: 每次训练的时间列表
"""
# 将日志列表转换为 NumPy 数组
recs = np.array(rec_loger)
pres = np.array(pre_loger)
map = np.array(map_loger)
mrr = np.array(mrr_loger)
fone = np.array(fone_loger)
# 找到召回率数组第一列的最大值
best_rec_0 = max(recs[:, 0])
# 获取最大值的索引
idx = list(recs[:, 0]).index(best_rec_0)
# 构建最终性能的字符串
final_perf = "Best Iter=[%d]@[%.1f]\trecall=[%s], precision=[%s], hit=[%s], ndcg=[%s], map=[%s],mrr=[%s], f1=[%s]" % \
(idx, time() - t0, '\t'.join(['%.5f' % r for r in recs[idx]]),
'\t'.join(['%.5f' % r for r in pres[idx]]),
'\t'.join(['%.5f' % r for r in map[idx]]),
'\t'.join(['%.5f' % r for r in mrr[idx]]),
'\t'.join(['%.5f' % r for r in fone[idx]]))
# 输出可被 Matlab 使用的结果
final_perf_output = "%s\n%s\n%s\n%s\n%s\n%s\n%s" % \
(','.join(['%.5f' % r for r in recs[idx]]),
','.join(['%.5f' % r for r in pres[idx]]),
','.join(['%.5f' % r for r in map[idx]]),
','.join(['%.5f' % r for r in mrr[idx]]),
','.join(['%.5f' % r for r in fone[idx]]))
print(final_perf) # 打印最终性能
# 计算平均训练时间并打印
avg_time = sum(training_time_list) / len(training_time_list)
time_consume = "Benchmarking time consuming: average {}s per epoch".format(avg_time)
print(time_consume)
# 定义结果保存的路径
results_path = '%soutput/%s/result.csv' % (args.proj_path, args.dataset)
# 确保目录存在
ensureDir(results_path)
f = open(results_path, 'a') # 以追加模式打开文件
f.write(final_perf_output) # 写入最终性能结果
f.close() # 关闭文件
# BPR 损失函数
def bpr_loss(self, users, pos_items, neg_items):
"""
计算 BPR (Bayesian Personalized Ranking) 损失
参数:
users: 表示用户的特征张量
pos_items: 表示正样本项目的特征张量
neg_items: 表示负样本项目的特征张量
返回:
mf_loss: 主要的损失值
emb_loss: 嵌入损失
reg_loss: 正则化损失
"""
# 计算用户和正样本项目的乘积,并在维度 1 上求和,得到正样本得分
pos_scores = torch.sum(torch.mul(users, pos_items), dim=1)
# 计算用户和负样本项目的乘积,并在维度 1 上求和,得到负样本得分
neg_scores = torch.sum(torch.mul(users, neg_items), dim=1)
# 计算正则化项
# 分别对用户、正样本项目和负样本项目的元素平方求和,然后乘以 1/2
# 最后将这三部分相加
regularizer = 1. / 2 * (users ** 2).sum() + 1. / 2 * (pos_items ** 2).sum() + 1. / 2 * (
neg_items ** 2).sum()
# 将正则化项除以批量大小进行归一化
regularizer = regularizer / self.batch_size
# 计算 pos_scores - neg_scores 的 log-sigmoid 值
maxi = F.logsigmoid(pos_scores - neg_scores)
# 计算 maxi 的负均值作为主要损失 mf_loss
mf_loss = -torch.mean(maxi)
# 计算嵌入损失 emb_loss,通过乘以衰减系数 self.decay 和正则化项
emb_loss = self.decay * regularizer
# 将正则化损失 reg_loss 设置为 0.0
reg_loss = 0.0
# 返回计算得到的三种损失
return mf_loss, emb_loss, reg_loss
def sparse_mx_to_torch_sparse_tensor(self, sparse_mx):
"""
将一个 Scipy 的稀疏矩阵转换为 PyTorch 的稀疏张量
参数:
sparse_mx: 输入的 Scipy 稀疏矩阵
返回:
转换后的 PyTorch 稀疏张量
"""
# 将输入的稀疏矩阵转换为 COO(坐标格式)格式,并转换数据类型为 np.float32
sparse_mx = sparse_mx.tocoo().astype(np.float32)
# 从转换后的 COO 格式的稀疏矩阵中提取行和列索引,并转换为 PyTorch 的张量,数据类型为 np.int64
indices = torch.from_numpy(
np.vstack((sparse_mx.row, sparse_mx.col)).astype(np.int64))
# 从转换后的 COO 格式的稀疏矩阵中提取数据,并转换为 PyTorch 的张量
values = torch.from_numpy(sparse_mx.data)
# 创建一个表示稀疏张量形状的 PyTorch 的 Size 对象
shape = torch.Size(sparse_mx.shape)
# 使用提取的索引、值和形状创建并返回 PyTorch 的稀疏张量
return torch.sparse.FloatTensor(indices, values, shape)
def get_sparse_tensor_value(self, X):
"""
从输入的稀疏矩阵中提取相关值
参数:
X: 输入的稀疏矩阵
返回:
indices: 行和列索引的矩阵
coo.data: 稀疏矩阵中的数据
coo.shape: 稀疏矩阵的形状
"""
# 将输入矩阵转换为 COO 格式,并转换数据类型为 np.float32
coo = X.tocoo().astype(np.float32)
# 构建行和列索引的矩阵
indices = np.mat([coo.row, coo.col]).transpose()
# 返回索引矩阵、数据和形状
return indices, coo.data, coo.shape
if __name__ == '__main__':
# 设置可见的 GPU 设备
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_id)
config = dict()
# 配置用户数量
config['n_users'] = data_generator.n_users
# 配置物品数量
config['n_items'] = data_generator.n_items
"""
*********************************************************
生成拉普拉斯矩阵,其中每个条目定义了两个连接节点之间的衰减因子(例如,p_ui)
"""
plain_adj, norm_adj, mean_adj = data_generator.get_adj_mat()
# 根据指定的邻接矩阵类型进行配置
if args.adj_type == 'norm':
config['norm_adj'] = norm_adj
print('use the normalized adjacency matrix')
else:
config['norm_adj'] = mean_adj + sp.eye(mean_adj.shape[0])
print('use the mean adjacency matrix')
t0 = time() # 记录开始时间
# 根据是否有预训练模型进行不同的操作
if args.pretrain == -1:
pretrain_data = load_pretrained_data()
else:
pretrain_data = None
# 创建模型包装器对象
Engine = Model_Wrapper(data_config=config, pretrain_data=pretrain_data)
if args.pretrain: # 如果使用预训练模型
print('pretrain path: ', Engine.weights_save_path)
if os.path.exists(Engine.weights_save_path): # 如果预训练模型路径存在
Engine.load_model() # 加载模型
users_to_test = list(data_generator.test_set.keys())
ret = Engine.test(users_to_test, drop_flag=True) # 进行测试
cur_best_pre_0 = ret['recall'][0]
# 打印预训练模型的评估指标结果
pretrain_ret = 'pretrained model recall=[%.5f, %.5f], precision=[%.5f, %.5f], hit=[%.5f, %.5f],' \
'ndcg=[%.5f, %.5f], map=[%.5f, %.5f], mrr=[%.5f, %.5f], f1=[%.5f, %.5f]' % \
(ret['recall'][0], ret['recall'][-1],
ret['precision'][0], ret['precision'][-1],
ret['hit_ratio'][0], ret['hit_ratio'][-1],
ret['ndcg'][0], ret['ndcg'][-1],
ret['map'][0], ret['map'][-1],
ret['mrr'][0], ret['mrr'][-1],
ret['fone'][0], ret['fone'][-1])
print(pretrain_ret)
else:
print('Cannot load pretrained model. Start training from stratch') # 无法加载预训练模型,从头开始训练
else:
print('without pretraining') # 不使用预训练
Engine.train() # 训练模型
if __name__ == '__main__':
# 设置可见的 GPU 设备
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_id)
config = dict()
# 配置用户数量
config['n_users'] = data_generator.n_users
# 配置物品数量
config['n_items'] = data_generator.n_items
"""
*********************************************************
生成拉普拉斯矩阵,其中每个条目定义了两个连接节点之间的衰减因子(例如,p_ui)
"""
plain_adj, norm_adj, mean_adj = data_generator.get_adj_mat()
# 根据指定的邻接矩阵类型进行配置
if args.adj_type == 'norm':
config['norm_adj'] = norm_adj
print('use the normalized adjacency matrix')
else:
config['norm_adj'] = mean_adj + sp.eye(mean_adj.shape[0])
print('use the mean adjacency matrix')
t0 = time() # 记录开始时间
# 根据是否有预训练模型进行不同的操作
if args.pretrain == -1:
pretrain_data = load_pretrained_data()
else:
pretrain_data = None
# 创建模型包装器对象
Engine = Model_Wrapper(data_config=config, pretrain_data=pretrain_data)
if args.pretrain: # 如果使用预训练模型
print('pretrain path: ', Engine.weights_save_path)
if os.path.exists(Engine.weights_save_path): # 如果预训练模型路径存在
Engine.load_model() # 加载模型
users_to_test = list(data_generator.test_set.keys())
ret = Engine.test(users_to_test, drop_flag=True) # 进行测试
cur_best_pre_0 = ret['recall'][0]
# 打印预训练模型的评估指标结果
pretrain_ret = 'pretrained model recall=[%.5f, %.5f], precision=[%.5f, %.5f], hit=[%.5f, %.5f],' \
'ndcg=[%.5f, %.5f], map=[%.5f, %.5f], mrr=[%.5f, %.5f], f1=[%.5f, %.5f]' % \
(ret['recall'][0], ret['recall'][-1],
ret['precision'][0], ret['precision'][-1],
ret['hit_ratio'][0], ret['hit_ratio'][-1],
ret['ndcg'][0], ret['ndcg'][-1],
ret['map'][0], ret['map'][-1],
ret['mrr'][0], ret['mrr'][-1],
ret['fone'][0], ret['fone'][-1])
print(pretrain_ret)
else:
print('Cannot load pretrained model. Start training from stratch') # 无法加载预训练模型,从头开始训练
else:
print('without pretraining') # 不使用预训练
Engine.train() # 训练模型