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)
sampling_mode='normal',
bern_flag=1,
filter_flag=1,
neg_ent=25,
neg_rel=0)
# dataloader for test
test_dataloader = TestDataLoader('./benchmarks/WN18RR/', 'link')
# define the model
simple = SimplE(
ent_tot=train_dataloader.get_ent_tot(),
rel_tot=train_dataloader.get_rel_tot(),
dim=200)
# define the loss function
model = NegativeSampling(
model=simple,
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=2000, alpha=0.5, use_gpu=True, opt_method='adagrad')
trainer.run()
simple.save_checkpoint('./checkpoint/simple.ckpt')
# test the model
simple.load_checkpoint('./checkpoint/simple.ckpt')
tester = Tester(model=simple, data_loader=test_dataloader, use_gpu=True)
tester.run_link_prediction(type_constrain=False)
# dataloader for training
train_dataloader = TrainDataLoader(
in_path="./benchmarks/FB15K237/",
nbatches=100,
threads=8,
sampling_mode="normal",
bern_flag=1,
filter_flag=1,
neg_ent=25,
neg_rel=0)
# dataloader for test
test_dataloader = TestDataLoader("./benchmarks/FB15K237/", "link")
# define the model
transd = TransD(ent_tot=train_dataloader.get_ent_tot(), rel_tot=train_dataloader.get_rel_tot(), dim_e=1, dim_r=1, p_norm=1, norm_flag=True)
# define the loss function
model = NegativeSampling(model=transd,loss=MarginLoss(margin=4.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()
transd.save_checkpoint('./checkpoint/transd_100.ckpt')
# test the model
transd.load_checkpoint('./checkpoint/transd_100.ckpt')
tester = Tester(model=transd, data_loader=test_dataloader, use_gpu=False)
tester.run_link_prediction(type_constrain=False)