diff --git a/dataloader/__init__.py b/common/__init__.py similarity index 100% rename from dataloader/__init__.py rename to common/__init__.py diff --git a/recommender/__init__.py b/common/config/__init__.py similarity index 100% rename from recommender/__init__.py rename to common/config/__init__.py diff --git a/common/config/parser.py b/common/config/parser.py new file mode 100644 index 0000000..962aab5 --- /dev/null +++ b/common/config/parser.py @@ -0,0 +1,100 @@ +import argparse + + +def parse_args(): + parser = argparse.ArgumentParser(description="Run KGPolicy2.") + # ------------------------- experimental settings specific for data set -------------------------------------------- + parser.add_argument( + "--data_path", nargs="?", default="../Data/", help="Input data path." + ) + parser.add_argument( + "--dataset", nargs="?", default="last-fm", help="Choose a dataset." + ) + parser.add_argument("--emb_size", type=int, default=64, help="Embedding size.") + parser.add_argument( + "--regs", + nargs="?", + default="1e-5", + help="Regularization for user and item embeddings.", + ) + parser.add_argument("--gpu_id", type=int, default=0, help="gpu id") + parser.add_argument( + "--k_neg", type=int, default=1, help="number of negative items in list" + ) + + # ------------------------- experimental settings specific for recommender ----------------------------------------- + parser.add_argument( + "--slr", type=float, default=0.0001, help="Learning rate for sampler." + ) + parser.add_argument( + "--rlr", type=float, default=0.0001, help="Learning rate recommender." + ) + + # ------------------------- experimental settings specific for sampler --------------------------------------------- + parser.add_argument( + "--edge_threshold", + type=int, + default=64, + help="edge threshold to filter knowledge graph", + ) + parser.add_argument( + "--num_sample", type=int, default=32, help="number fo samples from gcn" + ) + parser.add_argument( + "--k_step", type=int, default=2, help="k step from current positive items" + ) + parser.add_argument( + "--in_channel", type=str, default="[64, 32]", help="input channels for gcn" + ) + parser.add_argument( + "--out_channel", type=str, default="[32, 64]", help="output channels for gcn" + ) + parser.add_argument( + "--pretrain_s", + type=bool, + default=False, + help="load pretrained sampler data or not", + ) + + # ------------------------- experimental settings specific for training -------------------------------------------- + parser.add_argument( + "--batch_size", type=int, default=1024, help="batch size for training." + ) + parser.add_argument( + "--test_batch_size", type=int, default=1024, help="batch size for test" + ) + parser.add_argument("--num_threads", type=int, default=4, help="number of threads.") + parser.add_argument("--epoch", type=int, default=400, help="Number of epoch.") + parser.add_argument("--show_step", type=int, default=3, help="test step.") + parser.add_argument( + "--adj_epoch", type=int, default=1, help="build adj matrix per _ epoch" + ) + parser.add_argument( + "--pretrain_r", type=bool, default=True, help="use pretrained model or not" + ) + parser.add_argument( + "--freeze_s", + type=bool, + default=False, + help="freeze parameters of recommender or not", + ) + parser.add_argument( + "--model_path", + type=str, + default="model/best_fm.ckpt", + help="path for pretrain model", + ) + parser.add_argument( + "--out_dir", type=str, default="./weights/", help="output directory for model" + ) + parser.add_argument("--flag_step", type=int, default=32, help="early stop steps") + parser.add_argument( + "--gamma", type=float, default=0.99, help="gamma for reward accumulation" + ) + + # ------------------------- experimental settings specific for testing --------------------------------------------- + parser.add_argument( + "--Ks", nargs="?", default="[20, 40, 60, 80, 100]", help="evaluate K list" + ) + + return parser.parse_args() diff --git a/utility/__init__.py b/common/dataset/__init__.py similarity index 100% rename from utility/__init__.py rename to common/dataset/__init__.py diff --git a/common/dataset/build.py b/common/dataset/build.py new file mode 100644 index 0000000..9fffc27 --- /dev/null +++ b/common/dataset/build.py @@ -0,0 +1,22 @@ +from torch.utils.data import DataLoader +from common.dataset.dataset import TrainGenerator, TestGenerator + + +def build_loader(args_config, graph): + train_generator = TrainGenerator(args_config=args_config, graph=graph) + train_loader = DataLoader( + train_generator, + batch_size=args_config.batch_size, + shuffle=True, + num_workers=args_config.num_threads, + ) + + test_generator = TestGenerator(args_config=args_config, graph=graph) + test_loader = DataLoader( + test_generator, + batch_size=args_config.test_batch_size, + shuffle=False, + num_workers=args_config.num_threads, + ) + + return train_loader, test_loader diff --git a/dataloader/data_generator.py b/common/dataset/dataset.py similarity index 85% rename from dataloader/data_generator.py rename to common/dataset/dataset.py index 1246ab1..4f78c7a 100644 --- a/dataloader/data_generator.py +++ b/common/dataset/dataset.py @@ -26,7 +26,7 @@ def __getitem__(self, index): user_dict = self.user_dict # randomly select one user. u_id = random.sample(self.exist_users, 1)[0] - out_dict['u_id'] = u_id + out_dict["u_id"] = u_id # randomly select one positive item. pos_items = user_dict[u_id] @@ -35,16 +35,18 @@ def __getitem__(self, index): pos_idx = np.random.randint(low=0, high=n_pos_items, size=1)[0] pos_i_id = pos_items[pos_idx] - out_dict['pos_i_id'] = pos_i_id + out_dict["pos_i_id"] = pos_i_id neg_i_id = self.get_random_neg(pos_items, []) - out_dict['neg_i_id'] = neg_i_id + out_dict["neg_i_id"] = neg_i_id return out_dict def get_random_neg(self, pos_items, selected_items): while True: - neg_i_id = np.random.randint(low=self.low_item_index, high=self.high_item_index, size=1)[0] + neg_i_id = np.random.randint( + low=self.low_item_index, high=self.high_item_index, size=1 + )[0] if neg_i_id not in pos_items and neg_i_id not in selected_items: break @@ -63,6 +65,6 @@ def __getitem__(self, index): batch_data = {} u_id = self.users_to_test[index] - batch_data['u_id'] = u_id + batch_data["u_id"] = u_id return batch_data diff --git a/dataloader/data_processor.py b/common/dataset/preprocess.py similarity index 77% rename from dataloader/data_processor.py rename to common/dataset/preprocess.py index a2d0201..36c3dfe 100644 --- a/dataloader/data_processor.py +++ b/common/dataset/preprocess.py @@ -1,11 +1,7 @@ import collections import numpy as np import networkx as nx -import pickle -import os from tqdm import tqdm -from utility.helper import ensure_dir -from time import time class CFData(object): @@ -13,8 +9,8 @@ def __init__(self, args_config): self.args_config = args_config path = args_config.data_path + args_config.dataset - train_file = path + '/train.dat' - test_file = path + '/test.dat' + train_file = path + "/train.dat" + test_file = path + "/test.dat" # ----------get number of users and items & then load rating data from train_file & test_file------------ self.train_data = self._generate_interactions(train_file) @@ -30,10 +26,10 @@ def __init__(self, args_config): def _generate_interactions(file_name): inter_mat = list() - lines = open(file_name, 'r').readlines() + lines = open(file_name, "r").readlines() for l in lines: tmps = l.strip() - inters = [int(i) for i in tmps.split(' ')] + inters = [int(i) for i in tmps.split(" ")] u_id, pos_ids = inters[0], inters[1:] pos_ids = list(set(pos_ids)) @@ -70,19 +66,23 @@ def _id_range(train_mat, test_mat, idx): n_id = max_id - min_id + 1 return (min_id, max_id), n_id - self.user_range, self.n_users = _id_range(self.train_data, self.test_data, idx=0) - self.item_range, self.n_items = _id_range(self.train_data, self.test_data, idx=1) + self.user_range, self.n_users = _id_range( + self.train_data, self.test_data, idx=0 + ) + self.item_range, self.n_items = _id_range( + self.train_data, self.test_data, idx=1 + ) self.n_train = len(self.train_data) self.n_test = len(self.test_data) - print('-'*50) - print('- user_range: (%d, %d)' % (self.user_range[0], self.user_range[1])) - print('- item_range: (%d, %d)' % (self.item_range[0], self.item_range[1])) - print('- n_train: %d' % self.n_train) - print('- n_test: %d' % self.n_test) - print('- n_users: %d' % self.n_users) - print('- n_items: %d' % self.n_items) - print('-'*50) + print("-" * 50) + print("- user_range: (%d, %d)" % (self.user_range[0], self.user_range[1])) + print("- item_range: (%d, %d)" % (self.item_range[0], self.item_range[1])) + print("- n_train: %d" % self.n_train) + print("- n_test: %d" % self.n_test) + print("- n_users: %d" % self.n_users) + print("- n_items: %d" % self.n_items) + print("-" * 50) class KGData(object): @@ -92,7 +92,7 @@ def __init__(self, args_config, entity_start_id=0, relation_start_id=0): self.relation_start_id = relation_start_id path = args_config.data_path + args_config.dataset - kg_file = path + '/kg_final.txt' + kg_file = path + "/kg_final.txt" # ----------get number of entities and relations & then load kg data from kg_file ------------. self.kg_data, self.kg_dict, self.relation_dict = self._load_kg(kg_file) @@ -140,8 +140,8 @@ def _construct_kg(kg_np): def _statistic_kg_triples(self): def _id_range(kg_mat, idx): - min_id = min(min(kg_mat[:, idx]), min(kg_mat[:, 2-idx])) - max_id = max(max(kg_mat[:, idx]), max(kg_mat[:, 2-idx])) + min_id = min(min(kg_mat[:, idx]), min(kg_mat[:, 2 - idx])) + max_id = max(max(kg_mat[:, idx]), max(kg_mat[:, 2 - idx])) n_id = max_id - min_id + 1 return (min_id, max_id), n_id @@ -149,21 +149,31 @@ def _id_range(kg_mat, idx): self.relation_range, self.n_relations = _id_range(self.kg_data, idx=1) self.n_kg_triples = len(self.kg_data) - print('-'*50) - print('- entity_range: (%d, %d)' % (self.entity_range[0], self.entity_range[1])) - print('- relation_range: (%d, %d)' % (self.relation_range[0], self.relation_range[1])) - print('- n_entities: %d' % self.n_entities) - print('- n_relations: %d' % self.n_relations) - print('- n_kg_triples: %d' % self.n_kg_triples) - print('-'*50) + print("-" * 50) + print( + "- entity_range: (%d, %d)" % (self.entity_range[0], self.entity_range[1]) + ) + print( + "- relation_range: (%d, %d)" + % (self.relation_range[0], self.relation_range[1]) + ) + print("- n_entities: %d" % self.n_entities) + print("- n_relations: %d" % self.n_relations) + print("- n_kg_triples: %d" % self.n_kg_triples) + print("-" * 50) class CKGData(CFData, KGData): def __init__(self, args_config): CFData.__init__(self, args_config=args_config) - KGData.__init__(self, args_config=args_config, entity_start_id=self.n_users, relation_start_id=2) + KGData.__init__( + self, + args_config=args_config, + entity_start_id=self.n_users, + relation_start_id=2, + ) self.args_config = args_config - + self.ckg_graph = self._combine_cf_kg() def _combine_cf_kg(self): @@ -176,12 +186,12 @@ def _combine_cf_kg(self): # ... ids of other entities in range of [#users + #items, #users + #entities) # ... ids of relations in range of [0, 2 + 2 * #kg relations), including two 'interact' and 'interacted_by'. ckg_graph = nx.MultiDiGraph() - print('Begin to load interaction triples ...') + print("Begin to load interaction triples ...") for u_id, i_id in tqdm(cf_mat, ascii=True): ckg_graph.add_edges_from([(u_id, i_id)], r_id=0) ckg_graph.add_edges_from([(i_id, u_id)], r_id=1) - print('\nBegin to load knowledge graph triples ...') + print("\nBegin to load knowledge graph triples ...") for h_id, r_id, t_id in tqdm(kg_mat, ascii=True): ckg_graph.add_edges_from([(h_id, t_id)], r_id=r_id) return ckg_graph diff --git a/utility/test_model.py b/common/test.py similarity index 63% rename from utility/test_model.py rename to common/test.py index e372330..248a59b 100644 --- a/utility/test_model.py +++ b/common/test.py @@ -3,32 +3,32 @@ import numpy as np from tqdm import tqdm -from dataloader.data_loader import build_loader - def get_score(model, n_users, n_items, train_user_dict): u_e, i_e = torch.split(model.all_embed, [n_users, n_items]) - + score_matrix = torch.matmul(u_e, i_e.t()) for u, pos in train_user_dict.items(): - score_matrix[u][pos-n_users] = -1e5 - + score_matrix[u][pos - n_users] = -1e5 + return score_matrix + def cal_ndcg(topk, test_set, num_pos, k): n = min(num_pos, k) - nrange = np.arange(n)+2 - idcg = np.sum(1/np.log2(nrange)) + nrange = np.arange(n) + 2 + idcg = np.sum(1 / np.log2(nrange)) dcg = 0 for i, s in enumerate(topk): if s in test_set: - dcg += 1/np.log2(i+2) + dcg += 1 / np.log2(i + 2) - ndcg = dcg/idcg + ndcg = dcg / idcg return ndcg + def test_v2(model, ks, ckg): ks = eval(ks) train_user_dict, test_user_dict = ckg.train_user_dict, ckg.test_user_dict @@ -40,16 +40,18 @@ def test_v2(model, ks, ckg): score_matrix = get_score(model, n_users, n_items, train_user_dict) n_k = len(ks) - result = {'precision': np.zeros(n_k), - 'recall' : np.zeros(n_k), - 'ndcg': np.zeros(n_k), - 'hit_ratio': np.zeros(n_k)} - - for i, k in enumerate(tqdm(ks, ascii=True, desc='Evaluate')): - precision, recall, ndcg, hr = 0, 0, 0, 0 + result = { + "precision": np.zeros(n_k), + "recall": np.zeros(n_k), + "ndcg": np.zeros(n_k), + "hit_ratio": np.zeros(n_k), + } + + for i, k in enumerate(tqdm(ks, ascii=True, desc="Evaluate")): + precision, recall, ndcg, hr = 0, 0, 0, 0 _, topk_index = torch.topk(score_matrix, k) topk_index = topk_index.cpu().numpy() + n_users - + for test_u, gt_pos in test_user_dict.items(): topk = topk_index[test_u] num_pos = len(gt_pos) @@ -63,18 +65,10 @@ def test_v2(model, ks, ckg): hr += 1 if num_hit > 0 else 0 ndcg += cal_ndcg(topk, test_set, num_pos, k) - - result['precision'][i] = precision / n_test_users - result['recall'][i] = recall / n_test_users - result['ndcg'][i] = ndcg / n_test_users - result['hit_ratio'][i] = hr / n_test_users - return result + result["precision"][i] = precision / n_test_users + result["recall"][i] = recall / n_test_users + result["ndcg"][i] = ndcg / n_test_users + result["hit_ratio"][i] = hr / n_test_users - - - - - - - \ No newline at end of file + return result diff --git a/utility/helper.py b/common/utils.py similarity index 56% rename from utility/helper.py rename to common/utils.py index bcfaf15..b54033d 100644 --- a/utility/helper.py +++ b/common/utils.py @@ -17,37 +17,43 @@ def ensure_dir(dir_path): def uni2str(unicode_str): - return str(unicode_str.encode('ascii', 'ignore')).replace('\n', '').strip() + return str(unicode_str.encode("ascii", "ignore")).replace("\n", "").strip() def has_numbers(input_string): - return bool(re.search(r'\d', input_string)) + return bool(re.search(r"\d", input_string)) def del_multichar(input_string, chars): for ch in chars: - input_string = input_string.replace(ch, '') + input_string = input_string.replace(ch, "") return input_string def merge_two_dicts(x, y): - z = x.copy() # start with x's keys and values - z.update(y) # modifies z with y's keys and values & returns None + z = x.copy() # start with x's keys and values + z.update(y) # modifies z with y's keys and values & returns None return z -def early_stopping(log_value, best_value, stopping_step, expected_order='acc', flag_step=100): +def early_stopping( + log_value, best_value, stopping_step, expected_order="acc", flag_step=100 +): # early stopping strategy: - assert expected_order in ['acc', 'dec'] + assert expected_order in ["acc", "dec"] - if (expected_order == 'acc' and log_value >= best_value) or (expected_order == 'dec' and log_value <= best_value): + if (expected_order == "acc" and log_value >= best_value) or ( + expected_order == "dec" and log_value <= best_value + ): stopping_step = 0 best_value = log_value else: stopping_step += 1 if stopping_step >= flag_step: - print("Early stopping is trigger at step: {} log:{}".format(flag_step, log_value)) + print( + "Early stopping is trigger at step: {} log:{}".format(flag_step, log_value) + ) should_stop = True else: should_stop = False @@ -62,5 +68,5 @@ def freeze(model): def unfreeze(model): for param in model.parameters(): - param.requires_grad = True + param.requires_grad = True return model diff --git a/dataloader/data_loader.py b/dataloader/data_loader.py deleted file mode 100644 index 77f60ba..0000000 --- a/dataloader/data_loader.py +++ /dev/null @@ -1,14 +0,0 @@ -from torch.utils.data import DataLoader -from .data_generator import TrainGenerator, TestGenerator - - -def build_loader(args_config, graph): - train_generator = TrainGenerator(args_config=args_config, graph=graph) - train_loader = DataLoader(train_generator, batch_size=args_config.batch_size, shuffle=True, - num_workers=args_config.num_threads) - - test_generator = TestGenerator(args_config=args_config, graph=graph) - test_loader = DataLoader(test_generator, batch_size=args_config.test_batch_size, shuffle=False, - num_workers=args_config.num_threads) - - return train_loader, test_loader diff --git a/Main.py b/main.py similarity index 59% rename from Main.py rename to main.py index d43be87..689747e 100644 --- a/Main.py +++ b/main.py @@ -1,36 +1,37 @@ import os -import sys import random from time import time from pathlib import Path import torch -import torch.nn as nn import numpy as np from tqdm import tqdm from copy import deepcopy -import pickle -from utility.parser import parse_args -from utility.test_model import test_v2 -from utility.helper import early_stopping -from utility.parser import parse_args +from common.test import test_v2 +from common.utils import early_stopping +from common.config.parser import parse_args -from dataloader.data_loader import build_loader -from dataloader.data_processor import CKGData +from common.dataset.build import build_loader +from common.dataset.preprocess import CKGData -from recommender.MF import MF -from sampler.KGPolicy_Sampler import KGPolicy +from modules.recommender.MF import MF +from modules.sampler.kgpolicy import KGPolicy -def train_one_epoch(recommender, sampler, - train_loader, - recommender_optim, sampler_optim, - adj_matrix, edge_matrix, - train_data, - cur_epoch, - avg_reward): +def train_one_epoch( + recommender, + sampler, + train_loader, + recommender_optim, + sampler_optim, + adj_matrix, + edge_matrix, + train_data, + cur_epoch, + avg_reward, +): loss, base_loss, reg_loss = 0, 0, 0 epoch_reward = 0 @@ -39,8 +40,8 @@ def train_one_epoch(recommender, sampler, tbar = tqdm(train_loader, ascii=True) num_batch = len(train_loader) for batch_data in tbar: - - tbar.set_description('Epoch {}'.format(cur_epoch)) + + tbar.set_description("Epoch {}".format(cur_epoch)) if torch.cuda.is_available(): batch_data = {k: v.cuda(non_blocking=True) for k, v in batch_data.items()} @@ -56,7 +57,9 @@ def train_one_epoch(recommender, sampler, selected_neg_items = selected_neg_items_list[-1, :] train_set = train_data[users] - in_train = torch.sum(selected_neg_items.unsqueeze(1) == train_set.long(), dim=1).byte() + in_train = torch.sum( + selected_neg_items.unsqueeze(1) == train_set.long(), dim=1 + ).byte() selected_neg_items[in_train] = neg[in_train] base_loss_batch, reg_loss_batch = recommender(users, pos, selected_neg_items) @@ -67,8 +70,10 @@ def train_one_epoch(recommender, sampler, """Train sampler network""" sampler_optim.zero_grad() - selected_neg_items_list, selected_neg_prob_list = sampler(batch_data, adj_matrix, edge_matrix) - + selected_neg_items_list, selected_neg_prob_list = sampler( + batch_data, adj_matrix, edge_matrix + ) + with torch.no_grad(): reward_batch = recommender.get_reward(users, pos, selected_neg_items_list) @@ -89,24 +94,28 @@ def train_one_epoch(recommender, sampler, reinforce_loss = -1 * torch.sum(reward_batch * selected_neg_prob_list) reinforce_loss.backward() sampler_optim.step() - + """record loss in an epoch""" loss += loss_batch base_loss += base_loss_batch reg_loss += reg_loss_batch - + avg_reward = epoch_reward / num_batch - print(' Epoch {0:4d}: \ + print( + " Epoch {0:4d}: \ \n Training loss: [{1:4f} = {2:4f} + {3:4f}] \ - \n Reward: {4:4f}'.format(cur_epoch, loss, base_loss, reg_loss, avg_reward)) - + \n Reward: {4:4f}".format( + cur_epoch, loss, base_loss, reg_loss, avg_reward + ) + ) + return loss, base_loss, reg_loss, avg_reward def save_model(file_name, model, config): if not os.path.isdir(config.out_dir): os.mkdir(config.out_dir) - + model_file = Path(config.out_dir + file_name) model_file.touch(exist_ok=True) @@ -115,7 +124,7 @@ def save_model(file_name, model, config): def build_sampler_graph(n_nodes, edge_threshold, graph): - adj_matrix = torch.zeros(n_nodes, edge_threshold*2) + adj_matrix = torch.zeros(n_nodes, edge_threshold * 2) edge_matrix = torch.zeros(n_nodes, edge_threshold) """sample neighbors for each node""" @@ -125,14 +134,19 @@ def build_sampler_graph(n_nodes, edge_threshold, graph): sampled_edge = random.sample(neighbors, edge_threshold) edges = deepcopy(sampled_edge) else: - neg_id = random.sample(range(CKG.item_range[0], CKG.item_range[1]+1), edge_threshold-len(neighbors)) - node_id = [node]*(edge_threshold-len(neighbors)) + neg_id = random.sample( + range(CKG.item_range[0], CKG.item_range[1] + 1), + edge_threshold - len(neighbors), + ) + node_id = [node] * (edge_threshold - len(neighbors)) sampled_edge = neighbors + neg_id edges = neighbors + node_id - + """concatenate sampled edge with random edge""" - sampled_edge += random.sample(range(CKG.item_range[0], CKG.item_range[1]+1), edge_threshold) - + sampled_edge += random.sample( + range(CKG.item_range[0], CKG.item_range[1] + 1), edge_threshold + ) + adj_matrix[node] = torch.tensor(sampled_edge, dtype=torch.long) edge_matrix[node] = torch.tensor(edges, dtype=torch.long) @@ -163,11 +177,15 @@ def train(train_loader, test_loader, graph, data_config, args_config): train_data = build_train_data(train_mat) if args_config.pretrain_r: - print("\nLoad model from {}".format(args_config.data_path + args_config.model_path)) + print( + "\nLoad model from {}".format( + args_config.data_path + args_config.model_path + ) + ) paras = torch.load(args_config.data_path + args_config.model_path) all_embed = torch.cat((paras["user_para"], paras["item_para"])) data_config["all_embed"] = all_embed - + recommender = MF(data_config=data_config, args_config=args_config) sampler = KGPolicy(recommender, data_config, args_config) @@ -176,31 +194,36 @@ def train(train_loader, test_loader, graph, data_config, args_config): sampler = sampler.cuda() recommender = recommender.cuda() - print('\nSet sampler as: {}'.format(str(sampler))) - print('Set recommender as: {}\n'.format(str(recommender))) + print("\nSet sampler as: {}".format(str(sampler))) + print("Set recommender as: {}\n".format(str(recommender))) recommender_optimer = torch.optim.Adam(recommender.parameters(), lr=args_config.rlr) sampler_optimer = torch.optim.Adam(sampler.parameters(), lr=args_config.slr) loss_loger, pre_loger, rec_loger, ndcg_loger, hit_loger = [], [], [], [], [] - stopping_step, cur_best_pre_0, avg_reward = 0, 0., 0 + stopping_step, cur_best_pre_0, avg_reward = 0, 0.0, 0 t0 = time() for epoch in range(args_config.epoch): if epoch % args_config.adj_epoch == 0: """sample adjacency matrix""" - adj_matrix, edge_matrix = build_sampler_graph(data_config['n_nodes'], - args_config.edge_threshold, graph.ckg_graph) - - cur_epoch = epoch + 1 - loss, base_loss, reg_loss, avg_reward = train_one_epoch(recommender, sampler, - train_loader, - recommender_optimer, sampler_optimer, - adj_matrix, edge_matrix, - train_data, - cur_epoch, - avg_reward) + adj_matrix, edge_matrix = build_sampler_graph( + data_config["n_nodes"], args_config.edge_threshold, graph.ckg_graph + ) + cur_epoch = epoch + 1 + loss, base_loss, reg_loss, avg_reward = train_one_epoch( + recommender, + sampler, + train_loader, + recommender_optimer, + sampler_optimer, + adj_matrix, + edge_matrix, + train_data, + cur_epoch, + avg_reward, + ) """Test""" if cur_epoch % args_config.show_step == 0: @@ -210,30 +233,50 @@ def train(train_loader, test_loader, graph, data_config, args_config): t3 = 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']) + rec_loger.append(ret["recall"]) + pre_loger.append(ret["precision"]) + ndcg_loger.append(ret["ndcg"]) + hit_loger.append(ret["hit_ratio"]) - perf_str = 'Evaluate[%.1fs]: \ + perf_str = ( + "Evaluate[%.1fs]: \ \n recall=[%.5f, %.5f, %.5f, %.5f, %.5f], \ \n precision=[%.5f, %.5f, %.5f, %.5f, %.5f], \ \n hit=[%.5f, %.5f, %.5f, %.5f, %.5f], \ - \n ndcg=[%.5f, %.5f, %.5f, %.5f, %.5f] ' % \ - (t3 - t2, - ret['recall'][0], ret['recall'][1], - ret['recall'][2], ret['recall'][3], ret['recall'][4], - ret['precision'][0], ret['precision'][1], - ret['precision'][2], ret['precision'][3], ret['precision'][4], - ret['hit_ratio'][0], ret['hit_ratio'][1], - ret['hit_ratio'][2], ret['hit_ratio'][3], ret['hit_ratio'][4], - ret['ndcg'][0], ret['ndcg'][1], - ret['ndcg'][2], ret['ndcg'][3], ret['ndcg'][4]) + \n ndcg=[%.5f, %.5f, %.5f, %.5f, %.5f] " + % ( + t3 - t2, + ret["recall"][0], + ret["recall"][1], + ret["recall"][2], + ret["recall"][3], + ret["recall"][4], + ret["precision"][0], + ret["precision"][1], + ret["precision"][2], + ret["precision"][3], + ret["precision"][4], + ret["hit_ratio"][0], + ret["hit_ratio"][1], + ret["hit_ratio"][2], + ret["hit_ratio"][3], + ret["hit_ratio"][4], + ret["ndcg"][0], + ret["ndcg"][1], + ret["ndcg"][2], + ret["ndcg"][3], + ret["ndcg"][4], + ) + ) 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=args_config.flag_step) + cur_best_pre_0, stopping_step, should_stop = early_stopping( + ret["recall"][0], + cur_best_pre_0, + stopping_step, + expected_order="acc", + flag_step=args_config.flag_step, + ) if should_stop: break @@ -246,42 +289,52 @@ def train(train_loader, test_loader, graph, data_config, args_config): best_rec_0 = max(recs[:, 0]) idx = list(recs[:, 0]).index(best_rec_0) - final_perf = "Best Iter=[%d]@[%.1f]\n recall=[%s] \n precision=[%s] \n hit=[%s] \n ndcg=[%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 hit[idx]]), - '\t'.join(['%.5f' % r for r in ndcgs[idx]])) + final_perf = ( + "Best Iter=[%d]@[%.1f]\n recall=[%s] \n precision=[%s] \n hit=[%s] \n ndcg=[%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 hit[idx]]), + "\t".join(["%.5f" % r for r in ndcgs[idx]]), + ) + ) print(final_perf) -if __name__ == '__main__': +if __name__ == "__main__": """fix the random seed""" seed = 2020 random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) - + """initialize args and dataset""" args_config = parse_args() CKG = CKGData(args_config) - + """set the gpu id""" if torch.cuda.is_available(): torch.cuda.set_device(args_config.gpu_id) - - data_config = {'n_users': CKG.n_users, - 'n_items': CKG.n_items, - 'n_relations': CKG.n_relations + 2, - 'n_entities': CKG.n_entities, - 'n_nodes': CKG.entity_range[1] + 1, - 'item_range': CKG.item_range} - - print('\ncopying CKG graph for data_loader.. it might take a few minutes') + + data_config = { + "n_users": CKG.n_users, + "n_items": CKG.n_items, + "n_relations": CKG.n_relations + 2, + "n_entities": CKG.n_entities, + "n_nodes": CKG.entity_range[1] + 1, + "item_range": CKG.item_range, + } + + print("\ncopying CKG graph for data_loader.. it might take a few minutes") graph = deepcopy(CKG) train_loader, test_loader = build_loader(args_config=args_config, graph=graph) - - train(train_loader=train_loader, - test_loader=test_loader, - graph=CKG, - data_config=data_config, - args_config=args_config) + + train( + train_loader=train_loader, + test_loader=test_loader, + graph=CKG, + data_config=data_config, + args_config=args_config, + ) diff --git a/modules/__init__.py b/modules/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/recommender/KGAT.py b/modules/recommender/KGAT.py similarity index 76% rename from recommender/KGAT.py rename to modules/recommender/KGAT.py index a1b30db..de82726 100644 --- a/recommender/KGAT.py +++ b/modules/recommender/KGAT.py @@ -3,13 +3,15 @@ import torch.nn.functional as F import math -import torch_geometric as geometric +import torch_geometric as geometric + class GraphConv(nn.Module): """ Graph Convolutional Network embed CKG and using its embedding to calculate prediction score """ + def __init__(self, in_channel, out_channel): super(GraphConv, self).__init__() self.in_channel = in_channel @@ -23,14 +25,15 @@ def forward(self, x, edge_indices): x = F.normalize(x) return x + class KGAT(nn.Module): def __init__(self, data_config, args_config): super(KGAT, self).__init__() self.args_config = args_config self.data_config = data_config - self.n_users = data_config['n_users'] - self.n_items = data_config['n_items'] - self.n_nodes = data_config['n_nodes'] + self.n_users = data_config["n_users"] + self.n_items = data_config["n_items"] + self.n_nodes = data_config["n_nodes"] """set input and output channel manually""" input_channel = 64 @@ -43,9 +46,11 @@ def __init__(self, data_config, args_config): self.all_embed = self._init_weight() def _init_weight(self): - all_embed = nn.Parameter(torch.FloatTensor(self.n_nodes, self.emb_size), requires_grad=True) + all_embed = nn.Parameter( + torch.FloatTensor(self.n_nodes, self.emb_size), requires_grad=True + ) ui = self.n_users + self.n_items - + if self.args_config.pretrain_r: nn.init.xavier_uniform_(all_embed) all_embed.data[:ui] = self.data_config["all_embed"] @@ -53,27 +58,40 @@ def _init_weight(self): nn.init.xavier_uniform_(all_embed) return all_embed - + def build_edge(self, adj_matrix): """build edges based on adj_matrix""" sample_edge = self.args_config.edge_threshold edge_matrix = adj_matrix n_node = edge_matrix.size(0) - node_index = torch.arange(n_node, device=edge_matrix.device).unsqueeze(1).repeat(1, sample_edge).flatten() + node_index = ( + torch.arange(n_node, device=edge_matrix.device) + .unsqueeze(1) + .repeat(1, sample_edge) + .flatten() + ) neighbor_index = edge_matrix.flatten() - edges = torch.cat((node_index.unsqueeze(1), neighbor_index.unsqueeze(1)), dim=1) + edges = torch.cat((node_index.unsqueeze(1), neighbor_index.unsqueeze(1)), dim=1) return edges def forward(self, user, pos_item, neg_item, edges_matrix): - u_e, pos_e, neg_e = self.all_embed[user], self.all_embed[pos_item], self.all_embed[neg_item] + u_e, pos_e, neg_e = ( + self.all_embed[user], + self.all_embed[pos_item], + self.all_embed[neg_item], + ) edges = self.build_edge(edges_matrix) x = self.all_embed gcn_embedding = self.gcn(x, edges.t().contiguous()) - u_e_, pos_e_, neg_e_ = gcn_embedding[user], gcn_embedding[pos_item], gcn_embedding[neg_item] - + u_e_, pos_e_, neg_e_ = ( + gcn_embedding[user], + gcn_embedding[pos_item], + gcn_embedding[neg_item], + ) + u_e = torch.cat([u_e, u_e_], dim=1) pos_e = torch.cat([pos_e, pos_e_], dim=1) neg_e = torch.cat([neg_e, neg_e_], dim=1) @@ -100,7 +118,7 @@ def get_reward(self, users, pos_items, neg_items): neg_e = self.all_embed[neg_items] neg_scores = torch.sum(u_e * neg_e, dim=1) - ij = torch.sum(neg_e*pos_e, dim=1) + ij = torch.sum(neg_e * pos_e, dim=1) reward = neg_scores + ij return reward @@ -110,8 +128,10 @@ def _l2_loss(self, t): def inference(self, users): num_entity = self.n_nodes - self.n_users - self.n_items - user_embed, item_embed, _ = torch.split(self.all_embed, [self.n_users, self.n_items, num_entity], dim=0) - + user_embed, item_embed, _ = torch.split( + self.all_embed, [self.n_users, self.n_items, num_entity], dim=0 + ) + user_embed = user_embed[users] prediction = torch.matmul(user_embed, item_embed.t()) return prediction @@ -121,12 +141,12 @@ def rank(self, users, items): i_e = self.all_embed[items] u_e = u_e.unsqueeze(dim=1) - ranking = torch.sum(u_e*i_e, dim=2) + ranking = torch.sum(u_e * i_e, dim=2) ranking = ranking.squeeze() return ranking def __str__(self): - return "recommender using KGAT, embedding size {}".format(self.args_config.emb_size) - - + return "recommender using KGAT, embedding size {}".format( + self.args_config.emb_size + ) diff --git a/recommender/MF.py b/modules/recommender/MF.py similarity index 82% rename from recommender/MF.py rename to modules/recommender/MF.py index 7788e4b..016abf5 100644 --- a/recommender/MF.py +++ b/modules/recommender/MF.py @@ -9,8 +9,8 @@ def __init__(self, data_config, args_config): super(MF, self).__init__() self.args_config = args_config self.data_config = data_config - self.n_users = data_config['n_users'] - self.n_items = data_config['n_items'] + self.n_users = data_config["n_users"] + self.n_items = data_config["n_items"] self.emb_size = args_config.emb_size self.regs = eval(args_config.regs) @@ -18,8 +18,10 @@ def __init__(self, data_config, args_config): self.all_embed = self._init_weight() def _init_weight(self): - all_embed = nn.Parameter(torch.FloatTensor(self.n_users + self.n_items, self.emb_size)) - + all_embed = nn.Parameter( + torch.FloatTensor(self.n_users + self.n_items, self.emb_size) + ) + if self.args_config.pretrain_r: all_embed.data = self.data_config["all_embed"] else: @@ -61,7 +63,9 @@ def _l2_loss(t): return torch.sum(t ** 2) / 2 def inference(self, users): - user_embed, item_embed = torch.split(self.all_embed, [self.n_users, self.n_items], dim=0) + user_embed, item_embed = torch.split( + self.all_embed, [self.n_users, self.n_items], dim=0 + ) user_embed = user_embed[users] prediction = torch.matmul(user_embed, item_embed.t()) return prediction @@ -71,12 +75,12 @@ def rank(self, users, items): i_e = self.all_embed[items] u_e = u_e.unsqueeze(dim=1) - ranking = torch.sum(u_e*i_e, dim=2) + ranking = torch.sum(u_e * i_e, dim=2) ranking = ranking.squeeze() return ranking def __str__(self): - return "recommender using BPRMF, embedding size {}".format(self.args_config.emb_size) - - + return "recommender using BPRMF, embedding size {}".format( + self.args_config.emb_size + ) diff --git a/modules/recommender/__init__.py b/modules/recommender/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/modules/sampler/__init__.py b/modules/sampler/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/sampler/KGPolicy_Sampler.py b/modules/sampler/kgpolicy.py similarity index 83% rename from sampler/KGPolicy_Sampler.py rename to modules/sampler/kgpolicy.py index 2ea403d..e28b067 100644 --- a/sampler/KGPolicy_Sampler.py +++ b/modules/sampler/kgpolicy.py @@ -1,8 +1,8 @@ import torch import torch.nn as nn -import torch.nn.functional as F +import torch.nn.functional as F -import torch_geometric as geometric +import torch_geometric as geometric import networkx as nx from tqdm import tqdm @@ -14,10 +14,11 @@ class GraphConv(nn.Module): Input: embedding matrix for knowledge graph entity and adjacency matrix Output: gcn embedding for kg entity """ + def __init__(self, in_channel, out_channel, config): super(GraphConv, self).__init__() self.config = config - + self.conv1 = geometric.nn.SAGEConv(in_channel[0], out_channel[0]) self.conv2 = geometric.nn.SAGEConv(in_channel[1], out_channel[1]) @@ -29,7 +30,7 @@ def forward(self, x, edge_indices): x = self.conv2(x, edge_indices) x = F.dropout(x) x = F.normalize(x) - + return x @@ -39,6 +40,7 @@ class KGPolicy(nn.Module): Input: user, postive item, knowledge graph embedding Ouput: qualified negative item """ + def __init__(self, dis, params, config): super(KGPolicy, self).__init__() self.params = params @@ -52,7 +54,9 @@ def __init__(self, dis, params, config): self.n_entities = params["n_nodes"] self.item_range = params["item_range"] self.input_channel = in_channel - self.entity_embedding = self._initialize_weight(self.n_entities, self.input_channel) + self.entity_embedding = self._initialize_weight( + self.n_entities, self.input_channel + ) def _initialize_weight(self, n_entities, input_channel): """entities includes items and other entities in knowledge graph""" @@ -60,7 +64,9 @@ def _initialize_weight(self, n_entities, input_channel): kg_embedding = self.params["kg_embedding"] entity_embedding = nn.Parameter(kg_embedding) else: - entity_embedding = nn.Parameter(torch.FloatTensor(n_entities, input_channel[0])) + entity_embedding = nn.Parameter( + torch.FloatTensor(n_entities, input_channel[0]) + ) nn.init.xavier_uniform_(entity_embedding) if self.config.freeze_s: @@ -84,9 +90,13 @@ def forward(self, data_batch, adj_matrix, edge_matrix): """sample candidate negative items based on knowledge graph""" one_hop, one_hop_logits = self.kg_step(pos, users, adj_matrix, step=1) - candidate_neg, two_hop_logits = self.kg_step(one_hop, users, adj_matrix, step=2) + candidate_neg, two_hop_logits = self.kg_step( + one_hop, users, adj_matrix, step=2 + ) candidate_neg = self.filter_entity(candidate_neg, self.item_range) - good_neg, good_logits = self.prune_step(self.dis, candidate_neg, users, two_hop_logits) + good_neg, good_logits = self.prune_step( + self.dis, candidate_neg, users, two_hop_logits + ) good_logits = good_logits + one_hop_logits neg_list = torch.cat([neg_list, good_neg.unsqueeze(0)]) @@ -102,9 +112,14 @@ def build_edge(self, adj_matrix): edge_matrix = adj_matrix n_node = edge_matrix.size(0) - node_index = torch.arange(n_node, device=edge_matrix.device).unsqueeze(1).repeat(1, sample_edge).flatten() + node_index = ( + torch.arange(n_node, device=edge_matrix.device) + .unsqueeze(1) + .repeat(1, sample_edge) + .flatten() + ) neighbor_index = edge_matrix.flatten() - edges = torch.cat((node_index.unsqueeze(1), neighbor_index.unsqueeze(1)), dim=1) + edges = torch.cat((node_index.unsqueeze(1), neighbor_index.unsqueeze(1)), dim=1) return edges def kg_step(self, pos, user, adj_matrix, step): @@ -119,14 +134,14 @@ def kg_step(self, pos, user, adj_matrix, step): u_e = u_e.unsqueeze(dim=2) pos_e = gcn_embedding[pos] pos_e = pos_e.unsqueeze(dim=1) - + one_hop = adj_matrix[pos] i_e = gcn_embedding[one_hop] - p_entity = F.leaky_relu(pos_e * i_e) + p_entity = F.leaky_relu(pos_e * i_e) p = torch.matmul(p_entity, u_e) p = p.squeeze() - logits = F.softmax(p, dim=1) + logits = F.softmax(p, dim=1) """sample negative items based on logits""" batch_size = logits.size(0) @@ -162,7 +177,9 @@ def prune_step(dis, negs, users, logits): @staticmethod def filter_entity(neg, item_range): - random_neg = torch.randint(int(item_range[0]), int(item_range[1] + 1), neg.size(), device=neg.device) + random_neg = torch.randint( + int(item_range[0]), int(item_range[1] + 1), neg.size(), device=neg.device + ) neg[neg > item_range[1]] = random_neg[neg > item_range[1]] neg[neg < item_range[0]] = random_neg[neg < item_range[0]] diff --git a/utility/requirements.txt b/requirements.txt similarity index 100% rename from utility/requirements.txt rename to requirements.txt diff --git a/sampler/Adversarial_Sampler.py b/sampler/Adversarial_Sampler.py deleted file mode 100644 index 86531bc..0000000 --- a/sampler/Adversarial_Sampler.py +++ /dev/null @@ -1,108 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch.distributions import Categorical -import scipy.sparse as sp - -from utility.test_model import args_config, CKG - -class AdvNet(nn.Module): - """ - Adversarial Net: - Input: user-item interactions, i.e., pairs - Output: hard & informative policy (i.e., negative samples) with corresponding probability - """ - def __init__(self, data_config, args_config): - super(AdvNet, self).__init__() - - self.args_config = args_config - - self.n_users = data_config['n_users'] - self.n_items = data_config['n_items'] - - self.emb_size = args_config.emb_size - self.regs = args_config.regs - - self.policy_type = args_config.policy_type - - self.all_embed = self._init_weight() - self.sp_adj = self._generate_sp_adj() - - def _init_weight(self): - - all_embed = nn.Parameter(torch.FloatTensor(self.n_users + self.n_items, self.emb_size)) - nn.init.xavier_normal_(all_embed) - - return all_embed - - def forward(self, data_batch): - selected_neg_items, selected_neg_prob = self.neg_sampler(data_batch) - return selected_neg_items, selected_neg_prob - - def _generate_sp_adj(self): - train_data = CKG.train_data - rows = list(train_data[:, 0]) - cols = list(train_data[:, 1]) - vals = [1.] * len(rows) - - return sp.coo_matrix((vals, (rows, cols)), shape=(self.n_users + self.n_items, self.n_users + self.n_items)) - - def neg_sampler(self, data_batch): - def _get_sparse_train_mask(s, idx): - tmp = s.copy().tolil() - try: - tmp = tmp[idx.cpu().numpy()] - except Exception: - print(idx) - return tmp.tocoo() - - def _sparse_dense_mul(sp_coo, ds_torch): - rows = sp_coo.row - cols = sp_coo.col - vals = sp_coo.data - # get values from relevant entries of dense matrix - ds_vals = ds_torch[rows, cols] - return torch.sparse.FloatTensor(torch.LongTensor([rows, cols]), vals * ds_vals, ds_torch.size).to_dense() - - def _masking(ds_torch, train_mask): - rows = train_mask[0] - cols = train_mask[1] - - # get values from relevant entries of dense matrix - ds_vals = ds_torch[rows, cols] - return torch.sparse.FloatTensor(torch.LongTensor([rows, cols]), ds_vals, ds_torch.size).to_dense() - - user = data_batch['u_id'] - pos_item = data_batch['pos_i_id'] - train_mask = data_batch['train_mask'] - - u_e = self.all_embed[user] - pos_e = self.all_embed[pos_item] - all_e = self.all_embed - - # # get the mask for positive items appearing in the training set. - # sp_mask = _get_sparse_train_mask(self.sp_adj, idx=user) - - if self.policy_type == 'uj': - # ... (1) consider only user preference on item j; the larger, more likely to be selected. - policy_prob = torch.matmul(u_e, all_e.t()) - elif self.policy_type == 'uij': - # ... (2) consider user preference on item j, as well as similarity between item i and j. - policy_prob = torch.matmul(u_e + pos_e, all_e.t()) - else: - # ... (1) by default set as 'uij' - policy_prob = torch.matmul(u_e + pos_e, all_e.t()) - - # use softmax to calculate sampling probability. - # policy_prob = _sparse_dense_mul(sp_mask, policy_prob) - sp_mask = _masking(ds_torch=policy_prob, train_mask=train_mask) - policy_prob[policy_prob == 0] = -float("inf") - policy_prob = F.softmax(policy_prob, dim=1) - - # select one negative sample, based on the sampling probability. - policy_sampler = Categorical(policy_prob) - selected_neg_items = policy_sampler.sample() - raw_idx = range(u_e.size(0)) - selected_neg_prob = policy_prob[raw_idx, selected_neg_items] - - return selected_neg_items, selected_neg_prob \ No newline at end of file diff --git a/utility/parser.py b/utility/parser.py deleted file mode 100644 index 1e3741b..0000000 --- a/utility/parser.py +++ /dev/null @@ -1,70 +0,0 @@ -import argparse - - -def parse_args(): - parser = argparse.ArgumentParser(description="Run KGPolicy2.") - # ------------------------- experimental settings specific for data set -------------------------------------------- - parser.add_argument('--data_path', nargs='?', default='../Data/', - help='Input data path.') - parser.add_argument('--dataset', nargs='?', default='last-fm', - help='Choose a dataset.') - parser.add_argument('--emb_size', type=int, default=64, - help='Embedding size.') - parser.add_argument('--regs', nargs='?', default='1e-5', - help='Regularization for user and item embeddings.') - parser.add_argument('--gpu_id', type=int, default=0, - help='gpu id') - parser.add_argument('--k_neg', type=int, default=1, - help='number of negative items in list') - - # ------------------------- experimental settings specific for recommender ----------------------------------------- - parser.add_argument('--slr', type=float, default=0.0001, - help='Learning rate for sampler.') - parser.add_argument('--rlr', type=float, default=0.0001, - help='Learning rate recommender.') - - # ------------------------- experimental settings specific for sampler --------------------------------------------- - parser.add_argument('--edge_threshold', type=int, default=64, - help='edge threshold to filter knowledge graph') - parser.add_argument('--num_sample', type=int, default=32, - help='number fo samples from gcn') - parser.add_argument('--k_step', type=int, default=2, - help="k step from current positive items") - parser.add_argument('--in_channel', type=str, default='[64, 32]', - help='input channels for gcn') - parser.add_argument('--out_channel', type=str, default='[32, 64]', - help='output channels for gcn') - parser.add_argument('--pretrain_s', type=bool, default=False, - help="load pretrained sampler data or not") - - # ------------------------- experimental settings specific for training -------------------------------------------- - parser.add_argument('--batch_size', type=int, default=1024, - help='batch size for training.') - parser.add_argument('--test_batch_size', type=int, default=1024, - help='batch size for test') - parser.add_argument('--num_threads', type=int, default=4, - help='number of threads.') - parser.add_argument('--epoch', type=int, default=400, - help='Number of epoch.') - parser.add_argument('--show_step', type=int, default=3, - help='test step.') - parser.add_argument('--adj_epoch', type=int, default=1, - help='build adj matrix per _ epoch') - parser.add_argument('--pretrain_r', type=bool, default=True, - help="use pretrained model or not") - parser.add_argument('--freeze_s', type=bool, default=False, - help="freeze parameters of recommender or not") - parser.add_argument('--model_path', type=str, default='model/best_fm.ckpt', - help="path for pretrain model") - parser.add_argument("--out_dir", type=str, default='./weights/', - help='output directory for model') - parser.add_argument("--flag_step", type=int, default=32, - help="early stop steps") - parser.add_argument("--gamma", type=float, default=0.99, - help="gamma for reward accumulation") - - # ------------------------- experimental settings specific for testing --------------------------------------------- - parser.add_argument('--Ks', nargs='?', default='[20, 40, 60, 80, 100]', - help='evaluate K list') - - return parser.parse_args()