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
class Components:
from openke.module.model import DistMult
from openke.module.loss import SigmoidLoss
from openke.data import TrainDataLoader, TestDataLoader
from openke.module.strategy import NegativeSampling
from openke.config import Trainer, Tester
train_dataloader = TrainDataLoader
test_dataloader = TestDataLoader
model = DistMult
loss = SigmoidLoss(adv_temperature=0.5)
strategy = NegativeSampling
trainer = Trainer
tester = Tester
)
# 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=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')