diff --git a/openhgnn/config.ini b/openhgnn/config.ini index c2478071..7834c7d9 100644 --- a/openhgnn/config.ini +++ b/openhgnn/config.ini @@ -1,4 +1,33 @@ ; #################### add model config here +[MetaHIN] +; 修改前 +; input_dir=/home/zzh/Z测试代码/ +; output_dir=/home/zzh/Z测试代码/ + +input_dir=/openhgnn/dataset/Common_Dataset/ +output_dir=/openhgnn/dataset/Common_Dataset/ + +dataset=dbook +use_cuda= True +file_num= 10 +num_location= 453 +num_fea_item= 2 +num_publisher =1698 +num_fea_user= 1 +item_fea_len= 1 +embedding_dim= 32 +user_embedding_dim= 32 +item_embedding_dim= 32 +first_fc_hidden_dim= 64 +second_fc_hidden_dim= 64 +mp_update= 1 +local_update= 1 +lr= 5e-4 +mp_lr= 5e-3 +local_lr= 5e-3 +batch_size= 32 +num_epoch= 50 +seed=13 [FedHGNN] fea_dim = 64 diff --git a/openhgnn/config.py b/openhgnn/config.py index 951bf908..5725e4e6 100644 --- a/openhgnn/config.py +++ b/openhgnn/config.py @@ -41,7 +41,32 @@ def __init__(self, file_path, model, dataset, task, gpu): self.patience = conf.getint("General", "patience") self.mini_batch_flag = conf.getboolean("General", "mini_batch_flag") ############## add config.py ################# - + + elif self.model_name == "MetaHIN": + self.use_cuda = conf.getboolean("MetaHIN", "use_cuda") + self.file_num = conf.getint("MetaHIN", "file_num") + self.num_location = conf.getint("MetaHIN", "num_location") + self.num_fea_item = conf.getint("MetaHIN", "num_fea_item") + self.num_publisher = conf.getint("MetaHIN", "num_publisher") + self.num_fea_user = conf.getint("MetaHIN", "num_fea_user") + self.item_fea_len = conf.getint("MetaHIN", "item_fea_len") + self.embedding_dim = conf.getint("MetaHIN", "embedding_dim") + self.user_embedding_dim = conf.getint("MetaHIN", "user_embedding_dim") + self.item_embedding_dim = conf.getint("MetaHIN", "item_embedding_dim") + self.first_fc_hidden_dim = conf.getint("MetaHIN", "first_fc_hidden_dim") + self.second_fc_hidden_dim = conf.getint("MetaHIN", "second_fc_hidden_dim") + self.mp_update = conf.getint("MetaHIN", "mp_update") + self.local_update = conf.getint("MetaHIN", "local_update") + self.lr = conf.getfloat("MetaHIN", "lr") + self.mp_lr = conf.getfloat("MetaHIN", "mp_lr") + self.local_lr = conf.getfloat("MetaHIN", "local_lr") + self.batch_size = conf.getint("MetaHIN", "batch_size") + self.num_epoch = conf.getint("MetaHIN", "num_epoch") + self.input_dir = conf.get("MetaHIN", "input_dir") + self.output_dir = conf.get("MetaHIN", "output_dir") + self.seed = conf.getint("MetaHIN", "seed") + + elif self.model_name =='FedHGNN': self.fea_dim = conf.getint("FedHGNN","fea_dim") self.in_dim = conf.getint("FedHGNN","in_dim") diff --git a/openhgnn/dataset/MetaHIN_dataset.py b/openhgnn/dataset/MetaHIN_dataset.py new file mode 100644 index 00000000..81a66d39 --- /dev/null +++ b/openhgnn/dataset/MetaHIN_dataset.py @@ -0,0 +1,161 @@ + +import gc +import glob +import os +import pickle + +# from DataProcessor import Movielens +from tqdm import tqdm +from multiprocessing import Process, Pool +from multiprocessing.pool import ThreadPool +import numpy as np +import torch + + +class Meta_DataHelper: + def __init__(self, input_dir, config): + self.input_dir = input_dir + self.config = config + self.mp_list = ["ub", "ubab", "ubub"] + + from dgl.data.utils import download, extract_archive + # 只有dbook这一个数据集 + dataset_name = 'dbook' + self.zip_file = f'./openhgnn/dataset/Common_Dataset/{dataset_name}.zip' + # common_dataset/dbook_dir + self.base_dir = './openhgnn/dataset/Common_Dataset/' + dataset_name+'_dir' + self.url = f'https://s3.cn-north-1.amazonaws.com.cn/dgl-data/dataset/openhgnn/{dataset_name}.zip' + if os.path.exists(self.zip_file): + pass + else: + os.makedirs( os.path.join('./openhgnn/dataset/Common_Dataset/') ,exist_ok= True) + download(self.url, + path=os.path.join('./openhgnn/dataset/Common_Dataset/') + ) + if os.path.exists( self.base_dir ): + pass + else: + os.makedirs( os.path.join( self.base_dir ) ,exist_ok= True ) + extract_archive(self.zip_file, self.base_dir) + + + def load_data(self, data_set, state, load_from_file=True): + # 解压后的dbook目录： input_dir下的dbook文件夹 + # data_dir = self.input_dir + data_set + # 修改后代码 + data_dir = self.base_dir +'/'+data_set + supp_xs_s = [] + supp_ys_s = [] + supp_mps_s = [] + query_xs_s = [] + query_ys_s = [] + query_mps_s = [] + + if data_set == "yelp": + training_set_size = int( + len(glob.glob("{}/{}/*.npy".format(data_dir, state))) + / self.config.file_num + ) # support, query + + # load all data + for idx in tqdm(range(training_set_size)): + supp_xs_s.append( + torch.from_numpy( + np.load("{}/{}/support_x_{}.npy".format(data_dir, state, idx)) + ) + ) + supp_ys_s.append( + torch.from_numpy( + np.load("{}/{}/support_y_{}.npy".format(data_dir, state, idx)) + ) + ) + query_xs_s.append( + torch.from_numpy( + np.load("{}/{}/query_x_{}.npy".format(data_dir, state, idx)) + ) + ) + query_ys_s.append( + torch.from_numpy( + np.load("{}/{}/query_y_{}.npy".format(data_dir, state, idx)) + ) + ) + + supp_mp_data, query_mp_data = {}, {} + for mp in self.mp_list: + _cur_data = np.load( + "{}/{}/support_{}_{}.npy".format(data_dir, state, mp, idx), + encoding="latin1", + ) + supp_mp_data[mp] = [torch.from_numpy(x) for x in _cur_data] + _cur_data = np.load( + "{}/{}/query_{}_{}.npy".format(data_dir, state, mp, idx), + encoding="latin1", + ) + query_mp_data[mp] = [torch.from_numpy(x) for x in _cur_data] + supp_mps_s.append(supp_mp_data) + query_mps_s.append(query_mp_data) + else: # 'dbook' + training_set_size = int( + len(glob.glob("{}/{}/*.pkl".format(data_dir, state))) + / self.config.file_num + ) # support, query + + # load all data + for idx in tqdm(range(training_set_size)): + support_x = pickle.load( + open("{}/{}/support_x_{}.pkl".format(data_dir, state, idx), "rb") + ) + if support_x.shape[0] > 5: + continue + del support_x + supp_xs_s.append( + pickle.load( + open( + "{}/{}/support_x_{}.pkl".format(data_dir, state, idx), "rb" + ) + ) + ) + supp_ys_s.append( + pickle.load( + open( + "{}/{}/support_y_{}.pkl".format(data_dir, state, idx), "rb" + ) + ) + ) + query_xs_s.append( + pickle.load( + open("{}/{}/query_x_{}.pkl".format(data_dir, state, idx), "rb") + ) + ) + query_ys_s.append( + pickle.load( + open("{}/{}/query_y_{}.pkl".format(data_dir, state, idx), "rb") + ) + ) + + supp_mp_data, query_mp_data = {}, {} + for mp in self.mp_list: + supp_mp_data[mp] = pickle.load( + open( + "{}/{}/support_{}_{}.pkl".format(data_dir, state, mp, idx), + "rb", + ) + ) + query_mp_data[mp] = pickle.load( + open( + "{}/{}/query_{}_{}.pkl".format(data_dir, state, mp, idx), + "rb", + ) + ) + supp_mps_s.append(supp_mp_data) + query_mps_s.append(query_mp_data) + + print( + "#support set: {}, #query set: {}".format(len(supp_xs_s), len(query_xs_s)) + ) + total_data = list( + zip(supp_xs_s, supp_ys_s, supp_mps_s, query_xs_s, query_ys_s, query_mps_s) + ) # all training tasks + del (supp_xs_s, supp_ys_s, supp_mps_s, query_xs_s, query_ys_s, query_mps_s) + gc.collect() + return total_data \ No newline at end of file diff --git a/openhgnn/dataset/__init__.py b/openhgnn/dataset/__init__.py index b0ea2b7b..3f4ce971 100644 --- a/openhgnn/dataset/__init__.py +++ b/openhgnn/dataset/__init__.py @@ -16,6 +16,7 @@ from .SACN_dataset import * from .NBF_dataset import NBF_Dataset from .Ingram_dataset import Ingram_KG_TrainData, Ingram_KG_TestData +from .MetaHIN_dataset import Meta_DataHelper DATASET_REGISTRY = {} @@ -90,27 +91,17 @@ def build_dataset(dataset, task, *args, **kwargs): if isinstance(dataset, DGLDataset): return dataset +####### add dataset here if dataset == "meirec": train_dataloader = get_data_loader("train", batch_size=args[0]) test_dataloader = get_data_loader("test", batch_size=args[0]) return train_dataloader, test_dataloader - - - if dataset in CLASS_DATASETS: - return build_dataset_v2(dataset, task) - if not try_import_task_dataset(task): - exit(1) - - if dataset == 'NL-100': + elif dataset == 'NL-100': train_dataloader = Ingram_KG_TrainData('',dataset) valid_dataloader = Ingram_KG_TestData('', dataset,'valid') test_dataloader = Ingram_KG_TestData('',dataset,'test') return train_dataloader,valid_dataloader,test_dataloader - elif dataset == 'meirec': - train_dataloader = get_data_loader("train", batch_size=args[0]) - test_dataloader = get_data_loader("test", batch_size=args[0]) - return train_dataloader, test_dataloader elif dataset == 'AdapropT': dataload=AdapropTDataLoader(args) return dataload @@ -119,13 +110,22 @@ def build_dataset(dataset, task, *args, **kwargs): return dataload elif dataset == 'SACN' or dataset == 'LTE': return + elif dataset == "dbook": + dataload = Meta_DataHelper(args.input_dir, args) + return dataload + +############# + if dataset in CLASS_DATASETS: + return build_dataset_v2(dataset, task) + if not try_import_task_dataset(task): + exit(1) _dataset = None if dataset in ['aifb', 'mutag', 'bgs', 'am']: _dataset = 'rdf_' + task -##################### add dataset here +########### add dataset here elif dataset in ['acm4HGMAE','hgprompt_acm_dblp','acm4FedHGNN']: return DATASET_REGISTRY['common_dataset'](dataset, logger=kwargs['logger'],args = kwargs['args']) @@ -135,7 +135,7 @@ def build_dataset(dataset, task, *args, **kwargs): elif dataset in ['dblp4RHINE']: _dataset = 'rhine_'+task return DATASET_REGISTRY[_dataset](dataset, logger=kwargs['logger'],args = kwargs['args']) -###################### +########## elif dataset in ['acm4NSHE', 'acm4GTN', 'academic4HetGNN', 'acm_han', 'acm_han_raw', 'acm4HeCo', 'dblp', 'dblp4MAGNN', 'imdb4MAGNN', 'imdb4GTN', 'acm4NARS', 'demo_graph', 'yelp4HeGAN', 'DoubanMovie', diff --git a/openhgnn/models/MetaHIN.py b/openhgnn/models/MetaHIN.py new file mode 100644 index 00000000..098196c8 --- /dev/null +++ b/openhgnn/models/MetaHIN.py @@ -0,0 +1,532 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import dgl.nn as dglnn +from . import BaseModel, register_model +from torch.autograd import Variable +import numpy as np +from ..utils import Evaluator + + +@register_model("MetaHIN") +class MetaHIN(BaseModel): + @classmethod + def build_model_from_args(cls, args): + return cls(args, args.model_name) + + def __init__(self, config, model_name): + super(MetaHIN, self).__init__() + self.config = config + self.mp = ["ub", "ubab", "ubub"] + self.device = torch.device("cuda" if self.config.use_cuda else "cpu") + self.model_name = model_name + + self.item_emb = ItemEmbeddingDB(config) + self.user_emb = UserEmbeddingDB(config) + + self.mp_learner = MetapathLearner(config) + self.meta_learner = MetaLearner(config) + + self.mp_lr = config.mp_lr + self.local_lr = config.local_lr + self.emb_dim = self.config.embedding_dim + + self.cal_metrics = Evaluator(config.seed) + + self.ml_weight_len = len(self.meta_learner.update_parameters()) + self.ml_weight_name = list(self.meta_learner.update_parameters().keys()) + self.mp_weight_len = len(self.mp_learner.update_parameters()) + self.mp_weight_name = list(self.mp_learner.update_parameters().keys()) + + self.transformer_liners = self.transform_mp2task() + + self.meta_optimizer = torch.optim.Adam(self.parameters(), lr=config.lr) + + def transform_mp2task(self): + liners = {} + ml_parameters = self.meta_learner.update_parameters() + # output_dim_of_mp = self.config['user_embedding_dim'] + output_dim_of_mp = 32 # movielens: lr=0.001, avg mp, 0.8081 + for w in self.ml_weight_name: + liners[w.replace(".", "-")] = torch.nn.Linear( + output_dim_of_mp, np.prod(ml_parameters[w].shape) + ) + return torch.nn.ModuleDict(liners) + + def forward( + self, + support_user_emb, + support_item_emb, + support_set_y, + support_mp_user_emb, + vars_dict=None, + ): + """ """ + if vars_dict is None: + vars_dict = self.meta_learner.update_parameters() + + support_set_y_pred = self.meta_learner( + support_user_emb, support_item_emb, support_mp_user_emb, vars_dict + ) + loss = F.mse_loss(support_set_y_pred, support_set_y) + grad = torch.autograd.grad(loss, vars_dict.values(), create_graph=True) + + fast_weights = {} + for i, w in enumerate(vars_dict.keys()): + fast_weights[w] = vars_dict[w] - self.local_lr * grad[i] + + for idx in range( + 1, self.config.local_update + ): # for the current task, locally update + support_set_y_pred = self.meta_learner( + support_user_emb, + support_item_emb, + support_mp_user_emb, + vars_dict=fast_weights, + ) + loss = F.mse_loss( + support_set_y_pred, support_set_y + ) # calculate loss on support set + grad = torch.autograd.grad( + loss, fast_weights.values(), create_graph=True + ) # calculate gradients w.r.t. model parameters + + for i, w in enumerate(fast_weights.keys()): + fast_weights[w] = fast_weights[w] - self.local_lr * grad[i] + + return fast_weights + + def mp_update( + self, + support_set_x, + support_set_y, + support_set_mps, + query_set_x, + query_set_y, + query_set_mps, + ): + """ + Mete-update the parameters of MetaPathLearner, AggLearner and MetaLearner. + """ + # each mp + support_mp_enhanced_user_emb_s, query_mp_enhanced_user_emb_s = [], [] + mp_task_fast_weights_s = {} + mp_task_loss_s = {} + # 元路径学习器和元学习器（g与h）的初始权重 + mp_initial_weights = self.mp_learner.update_parameters() + ml_initial_weights = self.meta_learner.update_parameters() + # 提取出用户和物品的嵌入 + + support_user_emb = self.user_emb(support_set_x[:, self.config.item_fea_len :]) + support_item_emb = self.item_emb(support_set_x[:, 0 : self.config.item_fea_len]) + query_user_emb = self.user_emb(query_set_x[:, self.config.item_fea_len :]) + query_item_emb = self.item_emb(query_set_x[:, 0 : self.config.item_fea_len]) + # 对每一个元路径 + for mp in self.mp: + support_set_mp = list(support_set_mps[mp]) + query_set_mp = list(query_set_mps[mp]) + support_neighs_emb = self.item_emb(torch.cat(support_set_mp)) + support_index_list = list(map(lambda _: _.shape[0], support_set_mp)) + query_neighs_emb = self.item_emb(torch.cat(query_set_mp)) + query_index_list = list(map(lambda _: _.shape[0], query_set_mp)) + # 用元路径学习器计算用户增强的嵌入 + support_mp_enhanced_user_emb = self.mp_learner( # 对应论文中的聚合过程：g + support_user_emb, + support_item_emb, + support_neighs_emb, + mp, + support_index_list, + ) + # 用元学习器来预测 + support_set_y_pred = self.meta_learner( # 对应论文中的预测过程：h + support_user_emb, support_item_emb, support_mp_enhanced_user_emb + ) + # 损失和梯度 + loss = F.mse_loss(support_set_y_pred, support_set_y) + grad = torch.autograd.grad( + loss, mp_initial_weights.values(), create_graph=True + ) + # 更新mp的参数 + fast_weights = {} + for i in range(self.mp_weight_len): + weight_name = self.mp_weight_name[i] + fast_weights[weight_name] = ( + mp_initial_weights[weight_name] - self.mp_lr * grad[i] + ) + + # # 继续进行mp的元学习 + for idx in range(1, self.config.mp_update): + support_mp_enhanced_user_emb = self.mp_learner( + support_user_emb, + support_item_emb, + support_neighs_emb, + mp, + support_index_list, + vars_dict=fast_weights, + ) + support_set_y_pred = self.meta_learner( + support_user_emb, support_item_emb, support_mp_enhanced_user_emb + ) + loss = F.mse_loss(support_set_y_pred, support_set_y) + grad = torch.autograd.grad( + loss, fast_weights.values(), create_graph=True + ) + + for i in range(self.mp_weight_len): + weight_name = self.mp_weight_name[i] + fast_weights[weight_name] = ( + fast_weights[weight_name] - self.mp_lr * grad[i] + ) + ######################################################## + # 上面完成语义级适应,下面做任务级适应 + support_mp_enhanced_user_emb = self.mp_learner( + support_user_emb, + support_item_emb, + support_neighs_emb, + mp, + support_index_list, + vars_dict=fast_weights, + ) + support_mp_enhanced_user_emb_s.append(support_mp_enhanced_user_emb) + query_mp_enhanced_user_emb = self.mp_learner( + query_user_emb, + query_item_emb, + query_neighs_emb, + mp, + query_index_list, + vars_dict=fast_weights, + ) + query_mp_enhanced_user_emb_s.append(query_mp_enhanced_user_emb) + + f_fast_weights = {} + for w, liner in self.transformer_liners.items(): + w = w.replace("-", ".") + f_fast_weights[w] = ml_initial_weights[w] * torch.sigmoid( + liner(support_mp_enhanced_user_emb.mean(0)) + ).view(ml_initial_weights[w].shape) + # f_fast_weights = None + # # the current mp ---> task update + mp_task_fast_weights = self.forward( + support_user_emb, + support_item_emb, + support_set_y, + support_mp_enhanced_user_emb, + vars_dict=f_fast_weights, + ) + mp_task_fast_weights_s[mp] = mp_task_fast_weights + + query_set_y_pred = self.meta_learner( + query_user_emb, + query_item_emb, + query_mp_enhanced_user_emb, + vars_dict=mp_task_fast_weights, + ) + q_loss = F.mse_loss(query_set_y_pred, query_set_y) + mp_task_loss_s[mp] = q_loss.data # movielens: 0.8126 dbook 0.6084 + # mp_task_loss_s[mp] = loss.data # dbook 0.6144 + + # mp_att = torch.FloatTensor( + # [l / sum(mp_task_loss_s.values()) for l in mp_task_loss_s.values()] + # ).to( + # self.device + # ) # movielens: 0.81 + mp_att = F.softmax( + -torch.stack(list(mp_task_loss_s.values())), dim=0 + ) # movielens: 0.80781 lr0.001 + # mp_att = torch.FloatTensor([1.0 / len(self.config['mp'])] * len(self.config['mp'])).to(self.device) + + agg_task_fast_weights = self.aggregator(mp_task_fast_weights_s, mp_att) + agg_mp_emb = torch.stack(query_mp_enhanced_user_emb_s, 1) + # agg_mp_emb = torch.stack(support_mp_enhanced_user_emb_s, 1) + query_agg_enhanced_user_emb = torch.sum(agg_mp_emb * mp_att.unsqueeze(1), 1) + query_y_pred = self.meta_learner( + query_user_emb, + query_item_emb, + query_agg_enhanced_user_emb, + vars_dict=agg_task_fast_weights, + ) + + loss = F.mse_loss(query_y_pred, query_set_y) + query_y_real = query_set_y.data.cpu().numpy() + query_y_pred = query_y_pred.data.cpu().numpy() + mae, rmse = self.cal_metrics.prediction(query_y_real, query_y_pred) + ndcg_5 = self.cal_metrics.ranking(query_y_real, query_y_pred, k=5) + return loss, mae, rmse, ndcg_5 + + def global_update( + self, + support_xs, + support_ys, + support_mps, + query_xs, + query_ys, + query_mps, + device="cpu", + ): + """ """ + batch_sz = len(support_xs) + loss_s = [] + mae_s = [] + rmse_s = [] + ndcg_at_5_s = [] + + for i in range(batch_sz): # each task in a batch + support_mp = dict(support_mps[i]) # must be dict!!! + query_mp = dict(query_mps[i]) + + for mp in self.mp: + support_mp[mp] = map(lambda x: x.to(device), support_mp[mp]) + query_mp[mp] = map(lambda x: x.to(device), query_mp[mp]) + _loss, _mae, _rmse, _ndcg_5 = self.mp_update( + support_xs[i].to(device), + support_ys[i].to(device), + support_mp, + query_xs[i].to(device), + query_ys[i].to(device), + query_mp, + ) + loss_s.append(_loss) + mae_s.append(_mae) + rmse_s.append(_rmse) + ndcg_at_5_s.append(_ndcg_5) + + loss = torch.stack(loss_s).mean(0) + mae = np.mean(mae_s) + rmse = np.mean(rmse_s) + ndcg_at_5 = np.mean(ndcg_at_5_s) + + self.meta_optimizer.zero_grad() + loss.backward() + self.meta_optimizer.step() + + return loss.cpu().data.numpy(), mae, rmse, ndcg_at_5 + + def evaluation( + self, support_x, support_y, support_mp, query_x, query_y, query_mp, device="cpu" + ): + """ """ + support_mp = dict(support_mp) # must be dict!!! + query_mp = dict(query_mp) + for mp in self.mp: + support_mp[mp] = map(lambda x: x.to(device), support_mp[mp]) + query_mp[mp] = map(lambda x: x.to(device), query_mp[mp]) + + _, mae, rmse, ndcg_5 = self.mp_update( + support_x.to(device), + support_y.to(device), + support_mp, + query_x.to(device), + query_y.to(device), + query_mp, + ) + return mae, rmse, ndcg_5 + + def aggregator(self, task_weights_s, att): + for idx, mp in enumerate(self.mp): + if idx == 0: + att_task_weights = dict( + {k: v * att[idx] for k, v in task_weights_s[mp].items()} + ) + continue + tmp_att_task_weights = dict( + {k: v * att[idx] for k, v in task_weights_s[mp].items()} + ) + att_task_weights = dict( + zip( + att_task_weights.keys(), + list( + map( + lambda x: x[0] + x[1], + zip( + att_task_weights.values(), tmp_att_task_weights.values() + ), + ) + ), + ) + ) + + return att_task_weights + + def eval_no_MAML(self, query_set_x, query_set_y, query_set_mps): + # each mp + query_mp_enhanced_user_emb_s = [] + query_user_emb = self.user_emb(query_set_x[:, self.config.item_fea_len :]) + query_item_emb = self.item_emb(query_set_x[:, 0 : self.config.item_fea_len]) + + for mp in self.mp: + query_set_mp = list(query_set_mps[mp]) + query_neighs_emb = self.item_emb(torch.cat(query_set_mp)) + query_index_list = map(lambda _: _.shape[0], query_set_mp) + query_mp_enhanced_user_emb = self.mp_learner( + query_user_emb, query_item_emb, query_neighs_emb, mp, query_index_list + ) + query_mp_enhanced_user_emb_s.append(query_mp_enhanced_user_emb) + + mp_att = torch.FloatTensor([1.0 / len(self.mp)] * len(self.mp)).to( + self.device + ) # mean + agg_mp_emb = torch.stack(query_mp_enhanced_user_emb_s, 1) + query_agg_enhanced_user_emb = torch.sum(agg_mp_emb * mp_att.unsqueeze(1), 1) + + query_y_pred = self.meta_learner( + query_user_emb, query_item_emb, query_agg_enhanced_user_emb + ) + query_mae, query_rmse = self.cal_metrics.prediction( + query_set_y.data.cpu().numpy(), query_y_pred.data.cpu().numpy() + ) + query_ndcg_5 = self.cal_metrics.ranking( + query_set_y.data.cpu().numpy(), query_y_pred.data.cpu().numpy(), 5 + ) + + return query_mae, query_rmse, query_ndcg_5 + + +class MetaLearner(torch.nn.Module): + def __init__(self, config): + super(MetaLearner, self).__init__() + self.embedding_dim = config.embedding_dim + self.fc1_in_dim = 32 + config.item_embedding_dim + self.fc2_in_dim = config.first_fc_hidden_dim + self.fc2_out_dim = config.second_fc_hidden_dim + self.use_cuda = config.use_cuda + self.config = config + + # prediction parameters + self.vars = torch.nn.ParameterDict() + self.vars_bn = torch.nn.ParameterList() + + w1 = torch.nn.Parameter( + torch.ones([self.fc2_in_dim, self.fc1_in_dim]) + ) # 64, 96 + torch.nn.init.xavier_normal_(w1) + self.vars["ml_fc_w1"] = w1 + self.vars["ml_fc_b1"] = torch.nn.Parameter(torch.zeros(self.fc2_in_dim)) + + w2 = torch.nn.Parameter(torch.ones([self.fc2_out_dim, self.fc2_in_dim])) + torch.nn.init.xavier_normal_(w2) + self.vars["ml_fc_w2"] = w2 + self.vars["ml_fc_b2"] = torch.nn.Parameter(torch.zeros(self.fc2_in_dim)) + + w3 = torch.nn.Parameter(torch.ones([1, self.fc2_out_dim])) + torch.nn.init.xavier_normal_(w3) + self.vars["ml_fc_w3"] = w3 + self.vars["ml_fc_b3"] = torch.nn.Parameter(torch.zeros(1)) + + def forward(self, user_emb, item_emb, user_neigh_emb, vars_dict=None): + """ """ + if vars_dict is None: + vars_dict = self.vars + + x_i = item_emb + x_u = user_neigh_emb # movielens: loss:12.14... up! ; dbook 20epoch: user_cold: mae 0.6051; + + x = torch.cat((x_i, x_u), 1) # ?, item_emb_dim+user_emb_dim+user_emb_dim + x = F.relu(F.linear(x, vars_dict["ml_fc_w1"], vars_dict["ml_fc_b1"])) + x = F.relu(F.linear(x, vars_dict["ml_fc_w2"], vars_dict["ml_fc_b2"])) + x = F.linear(x, vars_dict["ml_fc_w3"], vars_dict["ml_fc_b3"]) + return x.squeeze() + + def zero_grad(self, vars_dict=None): + with torch.no_grad(): + if vars_dict is None: + for p in self.vars.values(): + if p.grad is not None: + p.grad.zero_() + else: + for p in vars_dict.values(): + if p.grad is not None: + p.grad.zero_() + + def update_parameters(self): + return self.vars + + +class MetapathLearner(torch.nn.Module): + def __init__(self, config): + super(MetapathLearner, self).__init__() + self.config = config + + # meta-path parameters + self.vars = torch.nn.ParameterDict() + neigh_w = torch.nn.Parameter( + torch.ones([32, config.item_embedding_dim]) + ) # dim=32, movielens 0.81006 + torch.nn.init.xavier_normal_(neigh_w) + self.vars["neigh_w"] = neigh_w + self.vars["neigh_b"] = torch.nn.Parameter(torch.zeros(32)) + + def forward(self, user_emb, item_emb, neighs_emb, mp, index_list, vars_dict=None): + """ """ + if vars_dict is None: + vars_dict = self.vars + agg_neighbor_emb = F.linear( + neighs_emb, vars_dict["neigh_w"], vars_dict["neigh_b"] + ) # (#neighbors, item_emb_dim) + output_emb = F.leaky_relu(torch.mean(agg_neighbor_emb, 0)).repeat( + user_emb.shape[0], 1 + ) # (#sample, user_emb_dim) + # + # # each mean, then att agg + # _emb = [] + # start = 0 + # for idx in index_list: + # end = start+idx + # _emb.append(F.leaky_relu(torch.mean(agg_neighbor_emb[start:end],0))) + # start = end + # output_emb = torch.stack(_emb, 0) # (#sample, dim) + + return output_emb + + def zero_grad(self, vars_dict=None): + with torch.no_grad(): + if vars_dict is None: + for p in self.vars.values(): + if p.grad is not None: + p.grad.zero_() + else: + for p in vars_dict.values(): + if p.grad is not None: + p.grad.zero_() + + def update_parameters(self): + return self.vars + + +class UserEmbeddingDB(torch.nn.Module): + def __init__(self, config): + super(UserEmbeddingDB, self).__init__() + self.num_location = config.num_location + self.embedding_dim = config.embedding_dim + + self.embedding_location = torch.nn.Embedding( + num_embeddings=self.num_location, embedding_dim=self.embedding_dim + ) + + def forward(self, user_fea): + """ + :param user_fea: tensor, shape = [#sample, #user_fea] + :return: + """ + location_idx = Variable(user_fea[:, 0], requires_grad=False) # [#sample] + location_emb = self.embedding_location(location_idx) + return location_emb # (1, 1*32) + + +class ItemEmbeddingDB(torch.nn.Module): + def __init__(self, config): + super(ItemEmbeddingDB, self).__init__() + self.num_publisher = config.num_publisher + self.embedding_dim = config.embedding_dim + + self.embedding_publisher = torch.nn.Embedding( + num_embeddings=self.num_publisher, embedding_dim=self.embedding_dim + ) + + def forward(self, item_fea): + """ + :param item_fea: + :return: + """ + publisher_idx = Variable(item_fea[:, 0], requires_grad=False) + publisher_emb = self.embedding_publisher(publisher_idx) # (1,32) + return publisher_emb # (1, 1*32) \ No newline at end of file diff --git a/openhgnn/models/__init__.py b/openhgnn/models/__init__.py index a1618a42..3099f58a 100644 --- a/openhgnn/models/__init__.py +++ b/openhgnn/models/__init__.py @@ -64,6 +64,7 @@ def build_model_from_args(args, hg): SUPPORTED_MODELS = { ##### add models here + "MetaHIN": "openhgnn.models.MetaHIN", 'HGA':'openhgnn.models.HGA', 'RHINE': 'openhgnn.models.RHINE', 'FedHGNN':'openhgnn.models.FedHGNN', diff --git a/openhgnn/tasks/__init__.py b/openhgnn/tasks/__init__.py index 86db3b05..01de3b4c 100644 --- a/openhgnn/tasks/__init__.py +++ b/openhgnn/tasks/__init__.py @@ -50,6 +50,7 @@ def try_import_task(task): SUPPORTED_TASKS = { + "coldstart_recommendation": "openhgnn.tasks.coldstart_recommendation", "KTN_trainer": "openhgnn.tasks.KTN", 'demo': 'openhgnn.tasks.demo', 'node_classification': 'openhgnn.tasks.node_classification', diff --git a/openhgnn/tasks/coldstart_recommendation.py b/openhgnn/tasks/coldstart_recommendation.py new file mode 100644 index 00000000..1079b387 --- /dev/null +++ b/openhgnn/tasks/coldstart_recommendation.py @@ -0,0 +1,32 @@ +import torch.nn.functional as F +import torch.nn as nn +from . import BaseTask, register_task +from ..dataset import build_dataset +from ..utils import Evaluator +import torch +import numpy as np +import dgl + + +@register_task("coldstart_recommendation") +class coldstart_recommendation(BaseTask): + def __init__(self, args): + super(coldstart_recommendation, self).__init__() + # self.logger = args.logger + self.dataloader = build_dataset( + args.dataset, "coldstart_recommendation", logger=args.logger, args=args + ) + self.args = args + self.evaluator = Evaluator(args.seed) + + def get_graph(self): + return + + def get_loss_fn(self): + return F.mse_loss + + def get_evaluator(self): + return self.evaluator.ndcg + + def evaluate(self, *args, **kwargs): + return \ No newline at end of file diff --git a/openhgnn/trainerflow/__init__.py b/openhgnn/trainerflow/__init__.py index 3f8f297f..95005ac5 100644 --- a/openhgnn/trainerflow/__init__.py +++ b/openhgnn/trainerflow/__init__.py @@ -48,6 +48,7 @@ def build_flow(args, flow_name): SUPPORTED_FLOWS = { + "coldstart_recommmendation": "openhgnn.trainerflow.coldstart_recommendation", 'SIAN_trainer': 'openhgnn.trainerflow.SIAN_trainer', 'entity_classification': 'openhgnn.trainerflow.entity_classification', 'node_classification': 'openhgnn.trainerflow.node_classification', diff --git a/openhgnn/trainerflow/base_flow.py b/openhgnn/trainerflow/base_flow.py index 23582409..667a7243 100644 --- a/openhgnn/trainerflow/base_flow.py +++ b/openhgnn/trainerflow/base_flow.py @@ -66,7 +66,8 @@ def __init__(self, args): self.max_epoch = args.max_epoch self.optimizer = None - if self.model_name in ["SIAN", "MeiREC", "ExpressGNN", "Ingram", "RedGNN","RedGNNT", "AdapropI", "AdapropT","RedGNNT", "Grail", "ComPILE","DisenKGAT"]: + if self.model_name in ["SIAN", "MeiREC", "ExpressGNN", "Ingram", "RedGNN","RedGNNT", "AdapropI", "AdapropT", + "RedGNNT", "Grail", "ComPILE","DisenKGAT","MetaHIN"]: return if self.model_name == "Ingram": return diff --git a/openhgnn/trainerflow/coldstart_recommendation.py b/openhgnn/trainerflow/coldstart_recommendation.py new file mode 100644 index 00000000..50012cf8 --- /dev/null +++ b/openhgnn/trainerflow/coldstart_recommendation.py @@ -0,0 +1,130 @@ +import torch +from tqdm import tqdm +from ..models import build_model +from . import BaseFlow, register_flow +import random +import time +from torch.utils.tensorboard import SummaryWriter +import numpy as np + +# python main.py -m MetaHIN -t coldstart_recommendation -d dbook -g 0 + +@register_flow("coldstart_recommendation") +class coldstart_recommendation(BaseFlow): + def __init__(self, args): + super(coldstart_recommendation, self).__init__(args) + self.datahelper = self.task.dataloader + self.states = [ + "meta_training", + "user_cold_testing", + "item_cold_testing", + "user_and_item_cold_testing", + "warm_up", + ] + self.data_set = "dbook" + self.model = ( + build_model(self.model).build_model_from_args(self.args).to(self.device) + ) + self.batch_size = args.batch_size + self.num_epoch = args.num_epoch + self.device = "cuda" + + def preprocess(self): + train_data = self.datahelper.load_data( + data_set=self.data_set, state="meta_training", load_from_file=True + ) + return train_data + + def full_test_step(self, model, device="cpu"): + print("evaluating model...") + if self.device != "cpu": + model.cuda() + model.eval() + for state in self.states: + if state == "meta_training": + continue + print(state + "...") + test_data = self.datahelper.load_data( + data_set=self.data_set, state=state, load_from_file=True + ) + supp_xs_s, supp_ys_s, supp_mps_s, query_xs_s, query_ys_s, query_mps_s = zip( + *test_data + ) # supp_um_s:(list,list,...,2553) + loss, mae, rmse = [], [], [] + ndcg_at_5 = [] + + for i in range(len(test_data)): # each task + _mae, _rmse, _ndcg_5 = model.evaluation( + supp_xs_s[i], + supp_ys_s[i], + supp_mps_s[i], + query_xs_s[i], + query_ys_s[i], + query_mps_s[i], + device, + ) + mae.append(_mae) + rmse.append(_rmse) + ndcg_at_5.append(_ndcg_5) + print( + " ndcg@5: {:.5f}".format( + np.mean(mae), np.mean(rmse), np.mean(ndcg_at_5) + ) + ) + + def train(self): + train_data = self.preprocess() + num_epoch = self.num_epoch + for i in range(num_epoch): # 20 + self.full_train_step(i, train_data) + if i % 10 == 0 and i != 0: + self.full_test_step(self.model, self.device) + self.model.train() + self.full_test_step(self.model, self.device) + + def full_train_step(self, epoch, train_data): + + batch_size = self.batch_size + + loss, mae, rmse = [], [], [] + ndcg_at_5 = [] + start = time.time() + random.shuffle(train_data) + num_batch = int(len(train_data) / batch_size) # ~80 + supp_xs_s, supp_ys_s, supp_mps_s, query_xs_s, query_ys_s, query_mps_s = zip( + *train_data + ) # supp_um_s:(list,list,...,2553) + for i in range( + num_batch + ): # each batch contains some tasks (each task contains a support set and a query set) + support_xs = list(supp_xs_s[batch_size * i : batch_size * (i + 1)]) + support_ys = list(supp_ys_s[batch_size * i : batch_size * (i + 1)]) + support_mps = list(supp_mps_s[batch_size * i : batch_size * (i + 1)]) + query_xs = list(query_xs_s[batch_size * i : batch_size * (i + 1)]) + query_ys = list(query_ys_s[batch_size * i : batch_size * (i + 1)]) + query_mps = list(query_mps_s[batch_size * i : batch_size * (i + 1)]) + + _loss, _mae, _rmse, _ndcg_5 = self.model.global_update( + support_xs, + support_ys, + support_mps, + query_xs, + query_ys, + query_mps, + self.device, + ) + loss.append(_loss) + mae.append(_mae) + rmse.append(_rmse) + ndcg_at_5.append(_ndcg_5) + + print( + "epoch: {}, loss: {:.6f}, cost time: {:.1f}s, mae: {:.5f}, rmse: {:.5f}, ndcg@5: {:.5f}".format( + epoch, + np.mean(loss), + time.time() - start, + np.mean(mae), + np.mean(rmse), + np.mean(ndcg_at_5), + ) + ) \ No newline at end of file diff --git a/openhgnn/utils/evaluator.py b/openhgnn/utils/evaluator.py index f44a58ac..db61127d 100644 --- a/openhgnn/utils/evaluator.py +++ b/openhgnn/utils/evaluator.py @@ -1,7 +1,7 @@ import numpy as np import torch as th from sklearn.cluster import KMeans -from sklearn.metrics import normalized_mutual_info_score, adjusted_rand_score +from sklearn.metrics import mean_absolute_error, mean_squared_error,normalized_mutual_info_score, adjusted_rand_score from sklearn.metrics import f1_score, accuracy_score, ndcg_score, roc_auc_score from sklearn.model_selection import train_test_split from sklearn.linear_model import LogisticRegression @@ -106,6 +106,38 @@ def ec_with_SVC(self, C, gamma, emd, labels, train_idx, test_idx): acc = metrics.accuracy_score(Y_test, Y_pred) return micro_f1, macro_f1, acc + + def prediction(self, real_score, pred_score): + MAE = mean_absolute_error(real_score, pred_score) + RMSE = math.sqrt(mean_squared_error(real_score, pred_score)) + return MAE, RMSE + + def dcg_at_k(self, scores): + # assert scores + return scores[0] + sum( + sc / math.log(ind + 1, 2) + for sc, ind in zip(scores[1:], range(2, len(scores) + 1)) + ) + + def ndcg_at_k(self, real_scores, predicted_scores): + idcg = self.dcg_at_k(sorted(real_scores, reverse=True)) + return (self.dcg_at_k(predicted_scores) / idcg) if idcg > 0.0 else 0.0 + + def ranking(self, real_score, pred_score, k): + # ndcg@k + sorted_idx = sorted( + np.argsort(real_score)[::-1][:k] + ) # get the index of the top k real score + r_s_at_k = real_score[sorted_idx] + p_s_at_k = pred_score[sorted_idx] + + ndcg_5 = self.ndcg_at_k(r_s_at_k, p_s_at_k) + + return ndcg_5 + + + + def filter(triplets_to_filter, target_s, target_r, target_o, num_entities, mode): triplets_to_filter = triplets_to_filter.copy() target_s, target_r, target_o = int(target_s), int(target_r), int(target_o)