def run_model(i, state):
print('newID:', i, state, len(state))
args.perf_file = os.path.join(directory, dataset + '_perf.txt')
torch.cuda.empty_cache()
# sleep to avoid multiple gpu occupy
time.sleep(10*(i%args.parrel)+1)
torch.cuda.set_device(select_gpu())
model = BaseModel(n_ent, n_rel, args, state)
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('ID:', i, 'structure:%s'%(str(state)), '\tvalid mrr', best_mrr)
for s in state:
f.write(str(s) + ' ')
f.write('\t\tbest_performance: '+best_str)
torch.cuda.empty_cache()
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}
space4kge = {
"lr": hp.uniform("lr", 0, 1),
"lamb": hp.uniform("lamb", -5, 0),
"decay_rate": hp.uniform("decay_rate", 0.99, 1.0),
"n_batch": hp.choice("n_batch", [128, 256, 512, 1024]),
"n_dim": hp.choice("n_dim", [64]),
}
trials = Trials()
best = fmin(run_kge,
space4kge,
algo=partial(tpe.suggest, n_startup_jobs=30),
max_evals=200,
trials=trials)
else:
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('\t\tbest_performance: ' + best_str + '\n')
