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')