def load_edge_data(self, load_geo_query=False, test_query_keep_graph=False): ''' just load 1-d query for train/val/test train_queries: train_queries[query_type][formula] = list of query val_queries: val_queries[one_neg/full_neg][query_type][formula] = list of query test_queries: test_queries[one_neg/full_neg][query_type][formula] = list of query ''' if load_geo_query: file_postfix = "-geo" else: file_postfix = "" print("Loading edge data..") print("Loading training edge data..") train_queries = load_queries_by_formula( self.args.data_dir + "/train_edges{:s}.pkl".format(file_postfix)) print("Loading validation edge data..") val_queries = load_test_queries_by_formula( self.args.data_dir + "/val_edges{:s}.pkl".format(file_postfix), keep_graph=test_query_keep_graph) print("Loading testing edge data..") test_queries = load_test_queries_by_formula( self.args.data_dir + "/test_edges{:s}.pkl".format(file_postfix), keep_graph=test_query_keep_graph) if load_geo_query: self.train_queries_geo = train_queries self.val_queries_geo = val_queries self.test_queries_geo = test_queries else: self.train_queries = train_queries self.val_queries = val_queries self.test_queries = test_queries
parser.add_argument("--log_dir", type=str, default="./log") parser.add_argument("--model_dir", type=str, default="./model") parser.add_argument("--decoder", type=str, default="bilinear") parser.add_argument("--inter_decoder", type=str, default="mean") parser.add_argument("--opt", type=str, default="adam") parser.add_argument("--pretrain", type=bool, default=False) args = parser.parse_args() print("Loading graph data..") graph, feature_modules, node_maps = load_graph(args.data_dir, args.embed_dim) if args.cuda: graph.features = cudify(feature_modules, node_maps) out_dims = {mode:args.embed_dim for mode in graph.relations} print("Loading edge data..") train_queries = load_queries_by_formula(args.data_dir + "/train_edges.pkl") val_queries = load_test_queries_by_formula(args.data_dir + "/val_edges.pkl") test_queries = load_test_queries_by_formula(args.data_dir + "/test_edges.pkl") print("Loading query data..") for i in range(2,4): train_queries.update(load_queries_by_formula(args.data_dir + "/train_queries_{:d}.pkl".format(i))) i_val_queries = load_test_queries_by_formula(args.data_dir + "/val_queries_{:d}.pkl".format(i)) val_queries["one_neg"].update(i_val_queries["one_neg"]) val_queries["full_neg"].update(i_val_queries["full_neg"]) i_test_queries = load_test_queries_by_formula(args.data_dir + "/test_queries_{:d}.pkl".format(i)) test_queries["one_neg"].update(i_test_queries["one_neg"]) test_queries["full_neg"].update(i_test_queries["full_neg"]) enc = get_encoder(args.depth, graph, out_dims, feature_modules, args.cuda, args.beta)
def load_multi_edge_query_data(self, load_geo_query=False, test_query_keep_graph=False): ''' Load multi edge query for train/val/test ''' if load_geo_query: file_postfix = "-geo" train_queries = self.train_queries val_queries = self.val_queries_geo test_queries = self.test_queries_geo else: file_postfix = "" train_queries = self.train_queries val_queries = self.val_queries test_queries = self.test_queries print("Loading {:s} query data..".format(file_postfix)) for i in range(2, 4): # if not args.kg_train: print("Loading training {:s} {:d} triple data..".format( file_postfix, i)) train_queries_file = self.args.data_dir + "/train_queries_{:d}{:s}.pkl".format( i, file_postfix) if path.exists(train_queries_file): train_queries.update( load_queries_by_formula(train_queries_file)) else: print("{} no exist!".format(train_queries_file)) print("Loading validate {:s} {:d} triple data..".format( file_postfix, i)) val_queries_file = self.args.data_dir + "/val_queries_{:d}{:s}.pkl".format( i, file_postfix) if path.exists(val_queries_file): i_val_queries = load_test_queries_by_formula( val_queries_file, keep_graph=test_query_keep_graph) val_queries["one_neg"].update(i_val_queries["one_neg"]) val_queries["full_neg"].update(i_val_queries["full_neg"]) else: print("{} no exist!".format(val_queries_file)) print("Loading testing {:s} {:d} triple data..".format( file_postfix, i)) test_queries_file = self.args.data_dir + "/test_queries_{:d}{:s}.pkl".format( i, file_postfix) if path.exists(test_queries_file): i_test_queries = load_test_queries_by_formula( test_queries_file, keep_graph=test_query_keep_graph) test_queries["one_neg"].update(i_test_queries["one_neg"]) test_queries["full_neg"].update(i_test_queries["full_neg"]) else: print("{} no exist!".format(test_queries_file)) if self.args.kg_train: print( "Loading x-inter train {:s} query data..".format(file_postfix)) if self.args.max_arity < 4: raise Exception("for full KG train, arity should be >= 4") for arity in range(4, self.args.max_arity + 1): print("Loading training {:s} {:d}-inter query data..".format( file_postfix, arity)) train_queries.update( load_queries_by_formula(self.args.data_dir + "/train_inter_queries_{:d}{:s}.pkl" .format(arity, file_postfix)))