def run_kge(params, struct=[0, 1, 2, 3]):
tester_val = lambda: model.test_link(valid_data, n_ent, heads, tails, args.
filter)
tester_tst = lambda: model.test_link(test_data, n_ent, heads, tails, args.
filter)
torch.cuda.set_device(select_gpu())
args.lr = params["lr"]
args.lamb = 10**params["lamb"]
args.decay_rate = params["decay_rate"]
args.n_batch = params["n_batch"]
args.n_dim = params["n_dim"]
plot_config(args)
model = BaseModel(n_ent, n_rel, args, struct)
best_mrr, best_str = model.train(train_data, corrupter, tester_val,
tester_tst)
with open(args.perf_file, 'a') as f:
print('structure:', struct)
print('write best mrr', best_str)
for s in struct:
f.write(str(s) + ' ')
#struc_str = ' '.join(struct) + '\t\t'
f.write('\t\tbest_performance: ' + best_str + '\n')
return best_mrr
def run_kge(params):
args.lr = params['lr']
args.lamb = 10**params['lamb']
args.decay_rate = params['decay_rate']
args.n_batch = params['n_batch']
args.n_dim = params['n_dim']
plot_config(args)
model = BaseModel(n_ent, n_rel, args, struct)
tester_val = lambda: model.test_link(valid_data, valid_head_filter,
valid_tail_filter)
tester_tst = lambda: model.test_link(test_data, test_head_filter,
test_tail_filter)
best_mrr, best_str = model.train(train_data, tester_val,
tester_tst)
with open(args.perf_file, 'a') as f:
print('structure:', struct, best_str)
for s in struct:
f.write(str(s) + ' ')
f.write(best_str + '\n')
return {'loss': -best_mrr, 'status': STATUS_OK}
def main(params, args, arch):
#def main(args, i):
# set number of threads in pytorch
torch.set_num_threads(6)
# select which gpu to use
logger_init(args)
# set gpu
if args.GPU:
torch.cuda.set_device(args.gpu)
# the default settings for correspdonding dataset
args = default_search_hyper(args)
# hyperOpt = {"lr":[0.00635456700742798, 0.0049700352658686425, 0.0023726642982752643],
# "lamb":[3.162503061522238e-05, 1.9567149674424395e-05, 1.0729611255307008e-05],
# "d":[512, 512, 512],
# "dr":[0.9933500551931267, 0.9903909316509071, 0.9933910046627364],
# "batch_size":[256, 256, 256]}
#
# args.lr = hyperOpt["lr"][i]
# args.lamb = hyperOpt["lamb"][i]
# args.n_dim = hyperOpt["d"][i]
# args.decay_rate = hyperOpt["dr"][i]
# args.n_batch = hyperOpt["batch_size"][i]
# load data
# read nary data
if args.n_arity > 2:
d = nary_dataloader(args.task_dir)
entity_idxs = {d.entities[i]: i for i in range(len(d.entities))}
relation_idxs = {d.relations[i]: i for i in range(len(d.relations))}
n_ent, n_rel = len(entity_idxs), len(relation_idxs)
print("Number of train:{}, valid:{}, test:{}.".format(
len(d.train_data), len(d.valid_data), len(d.test_data)))
train_data = torch.LongTensor(
get_data_idxs(d.train_data, entity_idxs, relation_idxs))
valid_data = torch.LongTensor(
get_data_idxs(d.valid_data, entity_idxs, relation_idxs))
test_data = torch.LongTensor(
get_data_idxs(d.test_data, entity_idxs, relation_idxs))
e1_sp, e2_sp, e3_sp = n_ary_heads(train_data, valid_data, test_data)
# train_data = torch.LongTensor(get_data_idxs(d.train_data, entity_idxs, relation_idxs))[0:512]
# valid_data = torch.LongTensor(get_data_idxs(d.valid_data, entity_idxs, relation_idxs))[0:512]
# test_data = torch.LongTensor(get_data_idxs(d.test_data, entity_idxs, relation_idxs))[0:512]
else:
loader = DataLoader(args.task_dir)
n_ent, n_rel = loader.graph_size()
train_data = loader.load_data('train')
valid_data = loader.load_data('valid')
test_data = loader.load_data('test')
print("Number of train:{}, valid:{}, test:{}.".format(
len(train_data[0]), len(valid_data[0]), len(test_data[0])))
heads, tails = loader.heads_tails()
train_data = torch.LongTensor(train_data).transpose(0, 1) #[0:512]
valid_data = torch.LongTensor(valid_data).transpose(0, 1) #[0:512]
test_data = torch.LongTensor(test_data).transpose(0, 1) #[0:512]
file_path = "search_nary" + "_" + str(args.num_blocks)
directory = os.path.join("results", args.dataset, file_path)
args.out_dir = directory
if not os.path.exists(directory):
os.makedirs(directory)
os.environ["OMP_NUM_THREADS"] = "4"
os.environ["MKL_NUM_THREADS"] = "4"
args.perf_file = os.path.join(
directory, args.dataset + '_search_nCP_nary_' + str(args.num_blocks) +
"_" + str(args.trial) + '.txt')
print('output file name:', args.perf_file)
args.lr = params["lr"]
args.decay_rate = params["decay_rate"]
args.n_batch = params["n_batch"]
#args.n_dim = params["n_dim"]
args.input_dropout = params["input_dropout"]
args.hidden_dropout = params["hidden_dropout"]
#args.lamb = params["lamb"]
plot_config(args)
def tester_val(facts=None, arch=None):
if args.n_arity == 2:
return model.test_link(test_data=valid_data,
n_ent=n_ent,
heads=heads,
tails=tails,
filt=args.filter,
arch=arch)
elif args.n_arity > 2:
return model.evaluate(valid_data, e1_sp, e2_sp, e3_sp, arch)
def tester_tst():
if args.n_arity == 2:
return model.test_link(test_data=test_data,
n_ent=n_ent,
heads=heads,
tails=tails,
filt=args.filter)
elif args.n_arity > 2:
return model.evaluate(test_data, e1_sp, e2_sp, e3_sp)
tester_trip_class = None
model = BaseModel(n_ent, n_rel, args, arch)
mrr = model.train(train_data, valid_data, tester_val, tester_tst,
tester_trip_class)
return mrr
logger_init(args)
task_dir = args.task_dir
loader = DataLoader(task_dir, args.N_1)
n_ent, n_rel = loader.graph_size()
train_data = loader.load_data('train')
valid_data = loader.load_data('valid')
test_data = loader.load_data('test')
args.n_train = len(train_data[0])
print("Number of train:{}, valid:{}, test:{}.".format(
len(train_data[0]), len(valid_data[0]), len(test_data[0])))
plot_config(args)
heads, tails = loader.heads_tails()
head_idx, tail_idx, head_cache, tail_cache, head_pos, tail_pos = loader.get_cache_list(
)
caches = [head_idx, tail_idx, head_cache, tail_cache, head_pos, tail_pos]
train_data = [torch.LongTensor(vec) for vec in train_data]
valid_data = [torch.LongTensor(vec) for vec in valid_data]
test_data = [torch.LongTensor(vec) for vec in test_data]
tester_val = lambda: model.test_link(valid_data, n_ent, heads, tails, args.
filter)
tester_tst = lambda: model.test_link(test_data, n_ent, heads, tails, args.
filter)
def train_stand(self, train_data, valid_data, derived_struct, rela_cluster, mrr):
self.rela_to_dict(rela_cluster)
self.args.perf_file = os.path.join(self.args.out_dir, self.args.dataset + '_std_' + str(self.args.m) + "_" + str(self.args.n) + "_" + str(mrr) + '.txt')
plot_config(self.args)
head, tail, rela = train_data
n_train = len(head)
if self.args.optim=='adam' or self.args.optim=='Adam':
self.optimizer = Adam(self.model.parameters(), lr=self.args.lr)
elif self.args.optim=='adagrad' or self.args.optim=='Adagrad':
self.optimizer = Adagrad(self.model.parameters(), lr=self.args.lr)
else:
self.optimizer = SGD(self.model.parameters(), lr=self.args.lr)
scheduler = ExponentialLR(self.optimizer, self.args.decay_rate)
n_batch = self.args.n_batch
start = time.time()
best_mrr = 0
for epoch in range(self.args.n_stand_epoch):
#self.epoch = epoch
rand_idx = torch.randperm(n_train)
if self.GPU:
head = head[rand_idx].cuda()
tail = tail[rand_idx].cuda()
rela = rela[rand_idx].cuda()
else:
head = head[rand_idx]
tail = tail[rand_idx]
rela = rela[rand_idx]
epoch_loss = 0
n_iters = 0
#lr = scheduler.get_lr()[0]
# train model weights
for h, t, r in batch_by_size(n_batch, head, tail, rela, n_sample=n_train):
self.model.zero_grad()
loss = self.model.forward(derived_struct, h, t, r, self.cluster_rela_dict)
loss += self.args.lamb * self.model.regul
loss.backward()
self.optimizer.step()
self.prox_operator()
epoch_loss += loss.data.cpu().numpy()
n_iters += 1
scheduler.step()
print("Epoch: %d/%d, Loss=%.2f, Stand Time=%.2f"%(epoch+1, self.args.n_stand_epoch, time.time()-start, epoch_loss/n_train))
if (epoch + 1) % 5 == 0:
test, randint = True, None
valid_mrr, valid_mr, valid_1, valid_3, valid_10 = self.tester_val(derived_struct, test, randint)
test_mrr, test_mr, test_1, test_3, test_10 = self.tester_tst(derived_struct, test, randint)
out_str = '%d \t %.2f \t %.2f \t %.4f %.1f %.4f %.4f %.4f\t%.4f %.1f %.4f %.4f %.4f\n' % (epoch, self.time_tot, epoch_loss/n_train,\
valid_mrr, valid_mr, valid_1, valid_3, valid_10, \
test_mrr, test_mr, test_1, test_3, test_10)
# output the best performance info
if test_mrr > best_mrr:
best_mrr = test_mrr
best_str = out_str
with open(self.args.perf_file, 'a+') as f:
f.write(out_str)
with open(self.args.perf_file, 'a+') as f:
f.write("best performance:" + best_str + "\n")
f.write("struct:" + str(derived_struct) + "\n")
f.write("rela:" + str(rela_cluster) + "\n")
return best_mrr
