n_user, n_item, n_entity, n_relation = data[0], data[1], data[2], data[3] # train_data: [user item score] train_data, test_data = data[4], data[5] # kg: [head relation tail] kg = data[6] mkr = MKR(args, n_user, n_item, n_entity, n_relation) loss_func = nn.BCELoss() optimizer_kge = optim.Adam(mkr.parameters(), lr=args.lr_kge) feed, tail_indices = get_data_for_kge(kg, 0, args.batch_size) feed, tail_indices = torch.Tensor(feed).long(), torch.Tensor(tail_indices).long() print("feed", feed) print("labels", tail_indices) for _ in range(10): tail_pred = mkr("kge", feed) tail_embeddings = mkr.entity_emb_matrix(tail_indices) for i in range(mkr.L): tail_embeddings = mkr.tail_mlps[i](tail_embeddings) # 定义均方根误差 rmse = (((tail_embeddings - tail_pred) ** 2).sum(1) / args.dim).sqrt().sum() print("rmse: {}".format(rmse.item())) optimizer_kge.zero_grad() rmse.backward() optimizer_kge.step()
def train(args, data): n_user, n_item, n_entity, n_relation = data[0], data[1], data[2], data[3] # train_data: [user item score] train_data, test_data = data[4], data[5] # kg: [head relation tail] kg = data[6] mkr = MKR(args, n_user, n_item, n_entity, n_relation) if args.cuda: mkr.cuda() loss_func = nn.BCELoss() optimizer_rs = optim.Adam(mkr.parameters(), lr=args.lr_rs) optimizer_kge = optim.Adam(mkr.parameters(), lr=args.lr_kge) # store best state best_test_acc = 0.0 best_state_dict = None for epoch in range(args.n_epochs): # RS training np.random.shuffle(train_data) start = 0 while start < train_data.shape[0]: feed, labels = get_data_for_rs(train_data, start, start + args.batch_size) feed, labels = torch.Tensor(feed).long(), torch.Tensor( labels).float() if args.cuda: feed, labels = feed.cuda(), labels.cuda() scores_normalized = mkr("rs", feed) # build loss for RS base_loss_rs = loss_func(scores_normalized, labels) l2_loss_rs = (mkr.user_embeddings** 2).sum() / 2 + (mkr.item_embeddings**2).sum() / 2 loss_rs = base_loss_rs + l2_loss_rs * args.l2_weight optimizer_rs.zero_grad() loss_rs.backward() optimizer_rs.step() start += args.batch_size if epoch % args.kge_interval == 0: # KGE training np.random.shuffle(kg) start = 0 while start < kg.shape[0]: feed, tail_indices = get_data_for_kge(kg, start, start + args.batch_size) feed, tail_indices = torch.Tensor(feed).long(), torch.Tensor( tail_indices).long() if args.cuda: feed, tail_indices = feed.cuda(), tail_indices.cuda() tail_pred = mkr("kge", feed) # build loss for KGE tail_embeddings = mkr.entity_emb_matrix(tail_indices) for i in range(mkr.L): tail_embeddings = mkr.tail_mlps[i](tail_embeddings) #scores_kge = mkr.sigmoid((tail_embeddings * tail_pred).sum(1)) #l2_loss_kge = (mkr.head_embeddings ** 2).sum() / 2 + (tail_embeddings ** 2).sum() / 2 #loss_kge = -scores_kge + l2_loss_kge * args.l2_weight rmse = (((tail_embeddings - tail_pred)**2).sum(1) / args.dim).sqrt().sum() optimizer_kge.zero_grad() rmse.backward() optimizer_kge.step() start += args.batch_size # Evaluating——train data inputs, y_true = get_data_for_rs(train_data, 0, train_data.shape[0]) inputs, y_true = torch.Tensor(inputs).long(), torch.Tensor( y_true).byte() if args.cuda: inputs, y_true = inputs.cuda(), y_true.cuda() train_auc, train_acc = mkr.evaluate(inputs, y_true) # Evaluating——test data inputs, y_true = get_data_for_rs(test_data, 0, test_data.shape[0]) inputs, y_true = torch.Tensor(inputs).long(), torch.Tensor( y_true).byte() if args.cuda: inputs, y_true = inputs.cuda(), y_true.cuda() test_auc, test_acc = mkr.evaluate(inputs, y_true) if test_acc > best_test_acc: best_test_acc = test_acc best_state_dict = deepcopy(mkr.state_dict()) print( "epoch {:3d} | train auc: {:.4f} acc: {:.4f} | test auc: {:.4f} acc:{:.4f}" .format(epoch, train_auc, train_acc, test_auc, test_acc)) # Save model wts_name = "../model/MKR_{}_{:.4f}.pth".format(args.dataset, best_test_acc) torch.save(best_state_dict, wts_name) print("Saved model to {}".format(wts_name)) return mkr