def train(args, data_info,
show_loss): # train 方法需要用到data info 是从 main 的 load data 方法里来的
train_data = data_info[0]
eval_data = data_info[1]
test_data = data_info[2]
n_entity = data_info[3]
n_relation = data_info[4]
ripple_set = data_info[5]
adj_entity = data_info[6]
adj_relation = data_info[7]
model = RippleNet(args, n_entity, n_relation, adj_entity,
adj_relation) # 在 train 里 用到model
if args.use_cuda:
model.cuda()
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
model.parameters()),
args.lr,
weight_decay=args.l2_weight)
for step in range(args.n_epoch):
# training
np.random.shuffle(train_data)
start = 0
while start < train_data.shape[0]:
return_dict = model(*get_feed_dict(args, model, train_data,
ripple_set, start, start +
args.batch_size))
loss = return_dict["loss"]
optimizer.zero_grad()
loss.backward()
optimizer.step()
start += args.batch_size
if show_loss:
print('%.1f%% %.4f' %
(start / train_data.shape[0] * 100, loss.item()))
# evaluation
train_auc, train_acc = evaluation(args, model, train_data, ripple_set,
args.batch_size)
eval_auc, eval_acc = evaluation(args, model, eval_data, ripple_set,
args.batch_size)
test_auc, test_acc = evaluation(args, model, test_data, ripple_set,
args.batch_size)
print(
'epoch %d train auc: %.4f acc: %.4f eval auc: %.4f acc: %.4f test auc: %.4f acc: %.4f'
% (step, train_auc, train_acc, eval_auc, eval_acc, test_auc,
test_acc))
def train(args, data_info, show_loss):
train_data = data_info[0]
eval_data = data_info[1]
test_data = data_info[2]
n_entity = data_info[3]
n_relation = data_info[4]
ripple_set = data_info[5]
model = RippleNet(args, n_entity, n_relation)
if args.use_cuda:
model.cuda()
optimizer = torch.optim.Adam(
filter(lambda p: p.requires_grad, model.parameters()),
args.lr,
)
for step in range(args.n_epoch):
# training
np.random.shuffle(train_data)
start = 0
while start < train_data.shape[0]:
return_dict = model(*get_feed_dict(args, model, train_data,
ripple_set, start, start +
args.batch_size))
loss = return_dict["loss"]
optimizer.zero_grad()
loss.backward()
optimizer.step()
start += args.batch_size
if show_loss:
print('%.1f%% %.4f' %
(start / train_data.shape[0] * 100, loss.item()))
# evaluation
# train_auc, train_acc, train_precision, train_recall, train_ndcg = evaluation(args, model, train_data, ripple_set, args.batch_size)
# eval_auc, eval_acc, eval_precision, eval_recall, eval_ndcg = evaluation(args, model, eval_data, ripple_set, args.batch_size)
test_precision, test_recall, test_ndcg = evaluation(
args, model, test_data, ripple_set, args.batch_size)
print('epoch %d test precision: %.4f recall: %.4f ndcg: %.4f' %
(step, test_precision, test_recall, test_ndcg))