train_data = read_data(os.path.join(task_dir, 'train.txt'), kb_index)
inplace_shuffle(*train_data)
valid_data = read_data(os.path.join(task_dir, 'valid.txt'), kb_index)
test_data = read_data(os.path.join(task_dir, 'test.txt'), kb_index)
heads, tails = heads_tails(n_ent, train_data, valid_data, test_data)
valid_data = [torch.LongTensor(vec) for vec in valid_data]
test_data = [torch.LongTensor(vec) for vec in test_data]
tester = lambda: gen.test_link(valid_data, n_ent, heads, tails)
train_data = [torch.LongTensor(vec) for vec in train_data]
mdl_type = config().pretrain_config
gen_config = config()[mdl_type]
if mdl_type == 'TransE':
corrupter = BernCorrupter(train_data, n_ent, n_rel)
gen = TransE(n_ent, n_rel, gen_config)
elif mdl_type == 'TransD':
corrupter = BernCorrupter(train_data, n_ent, n_rel)
gen = TransD(n_ent, n_rel, gen_config)
elif mdl_type == 'DistMult':
corrupter = BernCorrupterMulti(train_data, n_ent, n_rel,
gen_config.n_sample)
gen = DistMult(n_ent, n_rel, gen_config)
elif mdl_type == 'ComplEx':
corrupter = BernCorrupterMulti(train_data, n_ent, n_rel,
gen_config.n_sample)
gen = ComplEx(n_ent, n_rel, gen_config)
gen.pretrain(train_data, corrupter, tester)
gen.load(os.path.join(task_dir, gen_config.model_file))
gen.mdl.eval()
with torch.no_grad():
gen.test_link(test_data, n_ent, heads, tails)
gen.load(os.path.join(task_dir, gen_config.model_file))
dis.load(os.path.join(task_dir, dis_config.model_file))
train_data = read_data(os.path.join(task_dir, 'train.txt'), kb_index)
inplace_shuffle(*train_data)
valid_data = read_data(os.path.join(task_dir, 'valid.txt'), kb_index)
test_data = read_data(os.path.join(task_dir, 'test.txt'), kb_index)
filt_heads, filt_tails = heads_tails(n_ent, train_data, valid_data, test_data)
valid_data = [torch.LongTensor(vec) for vec in valid_data]
test_data = [torch.LongTensor(vec) for vec in test_data]
tester = lambda: dis.test_link(valid_data, n_ent, filt_heads, filt_tails)
train_data = [torch.LongTensor(vec) for vec in train_data]
dis.test_link(test_data, n_ent, filt_heads, filt_tails)
corrupter = BernCorrupterMulti(train_data, n_ent, n_rel, config().adv.n_sample)
src, rel, dst = train_data
n_train = len(src)
n_epoch = config().adv.n_epoch
n_batch = config().adv.n_batch
mdl_name = 'gan_dis_' + datetime.datetime.now().strftime("%m%d%H%M%S") + '.mdl'
best_perf = 0
avg_reward = 0
for epoch in range(n_epoch):
epoch_d_loss = 0
epoch_reward = 0
src_cand, rel_cand, dst_cand = corrupter.corrupt(src, rel, dst, keep_truth=False)
for s, r, t, ss, rs, ts in batch_by_num(n_batch, src, rel, dst, src_cand, rel_cand, dst_cand, n_sample=n_train):
gen_step = gen.gen_step(ss, rs, ts, temperature=config().adv.temperature)
src_smpl, dst_smpl = next(gen_step)
losses, rewards = dis.dis_step(s, r, t, src_smpl.squeeze(), dst_smpl.squeeze())