def train():
# dataloader for training
train_dataloader = TrainDataLoader(in_path='./data/kg/',
nbatches=100,
threads=8,
sampling_mode="normal",
bern_flag=1,
filter_flag=1,
neg_ent=25,
neg_rel=0)
# define the model
transe = TransE(ent_tot=train_dataloader.get_ent_tot(),
rel_tot=train_dataloader.get_rel_tot(),
dim=Config.entity_embedding_dim,
p_norm=1,
norm_flag=True)
# define the loss function
model = NegativeSampling(model=transe,
loss=MarginLoss(margin=5.0),
batch_size=train_dataloader.get_batch_size())
# train the model
trainer = Trainer(model=model,
data_loader=train_dataloader,
train_times=1000,
alpha=1.0,
use_gpu=True)
trainer.run()
transe.save_checkpoint('./data/kg/transe.ckpt')
def run_experiment(negative_rate, in_dim, out_dim, alpha):
train_dataloader = TrainDataLoader(in_path="./benchmarks/FB15K237/",
batch_size=10000,
threads=1,
sampling_mode="cross",
bern_flag=0,
filter_flag=1,
neg_ent=negative_rate,
neg_rel=0)
facts = TrainDataLoader(in_path="./benchmarks/FB15K237/",
batch_size=train_dataloader.get_triple_tot(),
threads=1,
sampling_mode="normal",
bern_flag=0,
filter_flag=1,
neg_ent=0,
neg_rel=0)
h, t, r, _, _ = [f for f in facts][0].values()
h = torch.Tensor(h).to(torch.long)
t = torch.Tensor(t).to(torch.long)
r = torch.Tensor(r).to(torch.long)
facts = torch.stack((h, r, t)).cuda().t()
test_dataloader = TestDataLoader("./benchmarks/FB15K237/", "link")
rotatte = RotAtte(
n_ent=train_dataloader.get_ent_tot(),
n_rel=train_dataloader.get_rel_tot(),
in_dim=in_dim,
out_dim=in_dim,
facts=facts,
negative_rate=negative_rate,
)
model = NegativeSampling(model=rotatte,
loss=SigmoidLoss(adv_temperature=2),
batch_size=train_dataloader.get_batch_size(),
regul_rate=0.0)
trainer = Trainer(model=model,
data_loader=train_dataloader,
train_times=100,
alpha=alpha,
use_gpu=True,
opt_method="adam")
trainer.run()
tester = Tester(model=rotatte, data_loader=test_dataloader, use_gpu=True)
result = tester.run_link_prediction(type_constrain=False)
MRR, MR, hits10, hits3, hits1 = result
return MRR
)
# dataloader for test
test_dataloader = TestDataLoader("./benchmarks/WN18RR/", "link")
# define the model
hole = HolE(
ent_tot = train_dataloader.get_ent_tot(),
rel_tot = train_dataloader.get_rel_tot(),
dim = 100
)
# define the loss function
model = NegativeSampling(
model = hole,
loss = SoftplusLoss(),
batch_size = train_dataloader.get_batch_size(),
regul_rate = 1.0
)
# train the model
trainer = Trainer(model = model, data_loader = train_dataloader, train_times = 1000, alpha = 0.5, use_gpu = True, opt_method = "adagrad")
trainer.run()
hole.save_checkpoint('./checkpoint/hole.ckpt')
# test the model
hole.load_checkpoint('./checkpoint/hole.ckpt')
tester = Tester(model = hole, data_loader = test_dataloader, use_gpu = True)
tester.run_link_prediction(type_constrain = False)
)
# dataloader for test
test_dataloader = TestDataLoader("./data/MetaQA/", "link")
# define the model
rotate = RotatE(
ent_tot = train_dataloader.get_ent_tot(),
rel_tot = train_dataloader.get_rel_tot(),
dim = 1024,
margin = 6.0,
epsilon = 2.0,
)
# define the loss function
model = NegativeSampling(
model = rotate,
loss = SigmoidLoss(adv_temperature = 2),
batch_size = train_dataloader.get_batch_size(),
regul_rate = 0.0
)
# train the model
trainer = Trainer(model = model, data_loader = train_dataloader, train_times = 1500, alpha = 2e-5, use_gpu = True, opt_method = "adam")
trainer.run()
rotate.save_checkpoint('./checkpoint/rotate.ckpt')
# test the model
rotate.load_checkpoint('./checkpoint/rotate.ckpt')
tester = Tester(model = rotate, data_loader = test_dataloader, use_gpu = True)
tester.run_link_prediction(type_constrain = False)
neg_ent=25,
neg_rel=0)
# dataloader for test
test_dataloader = TestDataLoader("./benchmarks/LUMB/", "link")
# define the model
transe = TransE(ent_tot=train_dataloader.get_ent_tot(),
rel_tot=train_dataloader.get_rel_tot(),
dim=200,
p_norm=1,
norm_flag=True)
# define the loss function
model = NegativeSampling(model=transe,
loss=MarginLoss(margin=5.0),
batch_size=train_dataloader.get_batch_size())
# train the model
trainer = Trainer(model=model,
data_loader=train_dataloader,
train_times=100,
alpha=1.0,
use_gpu=False)
trainer.run()
transe.save_checkpoint('./checkpoint/transe.ckpt')
# test the model
transe.load_checkpoint('./checkpoint/transe.ckpt')
tester = Tester(model=transe, data_loader=test_dataloader, use_gpu=False)
tester.run_link_prediction(type_constrain=False)
neg_ent=64,
neg_rel=0)
# dataloader for test
test_dataloader = TestDataLoader("./benchmarks/WN18RR/", "link")
# define the model
distmult = DistMult(ent_tot=train_dataloader.get_ent_tot(),
rel_tot=train_dataloader.get_rel_tot(),
dim=1024,
margin=200.0,
epsilon=2.0)
# define the loss function
model = NegativeSampling(model=distmult,
loss=SigmoidLoss(adv_temperature=0.5),
batch_size=train_dataloader.get_batch_size(),
l3_regul_rate=0.000005)
# train the model
trainer = Trainer(model=model,
data_loader=train_dataloader,
train_times=400,
alpha=0.002,
use_gpu=True,
opt_method="adam")
trainer.run()
distmult.save_checkpoint('./checkpoint/distmult.ckpt')
# test the model
distmult.load_checkpoint('./checkpoint/distmult.ckpt')
tester = Tester(model=distmult, data_loader=test_dataloader, use_gpu=True)
filter_flag = 1,
neg_ent = 25,
neg_rel = 0
)
# dataloader for test
test_dataloader = TestDataLoader("./benchmarks/FB15K237/", "link")
# define the model
rescal = RESCAL(
ent_tot = train_dataloader.get_ent_tot(),
rel_tot = train_dataloader.get_rel_tot(),
dim = 50
)
# define the loss function
model = NegativeSampling(
model = rescal,
loss = MarginLoss(margin = 1.0),
batch_size = train_dataloader.get_batch_size(),
)
# train the model
trainer = Trainer(model = model, data_loader = train_dataloader, train_times = 1000, alpha = 0.1, use_gpu = True, opt_method = "adagrad")
trainer.run()
rescal.save_checkpoint('./checkpoint/rescal.ckpt')
# test the model
rescal.load_checkpoint('./checkpoint/rescal.ckpt')
tester = Tester(model = rescal, data_loader = test_dataloader, use_gpu = True)
tester.run_link_prediction(type_constrain = False)
torch.from_numpy(
np.array(extract_path_vec_list).astype(dtype='float64')).float())
# define the model
transr = TransR(ent_tot=train_dataloader.get_ent_tot(),
rel_tot=train_dataloader.get_rel_tot(),
dim_e=30,
dim_r=50,
p_norm=1,
norm_flag=True,
rand_init=False)
transr.load_rel_embeddings(rel_embedding)
model_r = NegativeSampling(model=transr,
loss=MarginLoss(margin=3.0),
batch_size=train_dataloader.get_batch_size())
for k, v in model_r.named_parameters():
if k == 'model.rel_embeddings.weight':
v.requires_grad = False
# train transr
# transr.set_parameters(parameters)
trainer = Trainer(model=model_r,
data_loader=train_dataloader,
train_times=200,
alpha=1.0,
use_gpu=False)
trainer.run()
transr.save_checkpoint('./checkpoint/fault_dataset_transr.ckpt')
'entity': '../concept_glove.max.npy',
'relation': '../relation_glove.max.npy'
}
# define the model
transe = TransE(ent_tot=train_dataloader.get_ent_tot(),
rel_tot=train_dataloader.get_rel_tot(),
dim=100,
p_norm=1,
margin=1.0,
norm_flag=True,
init='pretrain',
init_weights=pretrain_init)
# define the loss function
model = NegativeSampling(model=transe,
loss=SigmoidLoss(adv_temperature=1),
batch_size=train_dataloader.get_batch_size())
# train the model
checkpoint_dir = Path('./checkpoint/')
checkpoint_dir.mkdir(exist_ok=True, parents=True)
alpha = 0.001
trainer = Trainer(model=model,
data_loader=train_dataloader,
train_times=1000,
alpha=alpha,
use_gpu=torch.cuda.is_available(),
opt_method='adam')
trainer.run()
transe.save_checkpoint(f'./checkpoint/transe_{transe.init}.ckpt')