コード例 #1
0
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
コード例 #2
0
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
コード例 #3
0
ファイル: train_rotate.py プロジェクト: yeeeqichen/KGQA
)

# 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)
コード例 #4
0
                                   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')