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')
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=25, neg_rel=0) # dataloader for test test_dataloader = TestDataLoader("./data/MetaQA/", "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=1000, alpha=1.0, use_gpu=True) 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=True) tester.run_link_prediction(type_constrain=False)
threads=8, sampling_mode="normal", bern_flag=1, filter_flag=1, neg_ent=25, neg_rel=5) # define the model transe = TransE(ent_tot=train_dataloader.get_ent_tot(), rel_tot=train_dataloader.get_rel_tot(), dim=200, p_norm=2, 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() if not os.path.exists(os.path.join('checkpoint', phase)): os.mkdir(os.path.join('checkpoint', phase)) save_path = os.path.join('checkpoint', phase, 'transe_pso.ckpt') transe.save_checkpoint(save_path)
for k, v in model.named_parameters(): if k == 'model.rel_embeddings.weight': v.requires_grad = False # train the model trainer = Trainer(model=model, data_loader=train_dataloader, train_times=1000, alpha=1.0, use_gpu=False) trainer.run() relationship = trainer.model.model.rel_embeddings.weight.data entity = trainer.model.model.ent_embeddings.weight.data relationship = trainer.model.model.rel_embeddings.weight.data entity_np = np.array(entity) np.savetxt('./embedding/TransE/fault_dataset_entity_result.txt', entity_np) relationship_np = np.array(relationship) np.savetxt('./embedding/TransE/fault_dataset_relationship_result.txt', relationship_np) transe.save_checkpoint('./checkpoint/fault_dataset_transe.ckpt') # test the model transe.load_checkpoint('./checkpoint/fault_dataset_transe.ckpt') tester = Tester(model=transe, data_loader=test_dataloader, use_gpu=False) tester.run_link_prediction(type_constrain=False)
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') # test the model #transe.load_checkpoint('./checkpoint/transe.ckpt') #tester = Tester(model = transe, data_loader = test_dataloader, use_gpu = torch.cuda.is_available()) #tester.run_link_prediction(type_constrain = False)
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=1, p_norm=1, norm_flag=False) # 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_100_lumb.ckpt') # test the model transe.load_checkpoint('./checkpoint/transe_100_lumb.ckpt') #tester = Tester(model=transe, data_loader=test_dataloader, use_gpu=False) #tester.run_link_prediction(type_constrain=False)
dim=1024, p_norm=1, norm_flag=True) # define the loss function model = NegativeSampling(model=transe, loss=MarginLoss(margin=0.0), batch_size=train_dataloader.get_batch_size()) # train the model trainer = Trainer(model=model, data_loader=train_dataloader, train_times=1000, alpha=1, use_gpu=True) trainer.run() transe.save_checkpoint('./checkpoint/CommonGen_transe.ckpt') # test the model transe.load_checkpoint('./checkpoint/CommonGen_transe.ckpt') tester = Tester(model=transe, data_loader=test_dataloader, use_gpu=True) tester.run_link_prediction(type_constrain=False) import pickle import numpy as np ent_embeddings = transe.ent_embeddings.weight.data.cpu().detach().numpy() rel_embeddings = transe.rel_embeddings.weight.data.cpu().detach().numpy() pickle.dump(ent_embeddings, open("./benchmarks/CommonGen/ent_embeddings", "wb")) pickle.dump(rel_embeddings, open("./benchmarks/CommonGen/rel_embeddings", "wb"))
type_constrain=False) # define the model transe = TransE(ent_tot=train_dataloader.get_ent_tot(), rel_tot=train_dataloader.get_rel_tot(), dim=100, p_norm=1, norm_flag=True) # define the loss function model = NegativeSampling(model=transe, loss=MarginLoss(margin=2.0), batch_size=train_dataloader.get_batch_size()) # train the model trainer = Trainer(model=model, data_loader=train_dataloader, train_times=epochs, alpha=alpha, use_gpu=True) trainer.run() transe.save_checkpoint(os.path.join(ckpt_path, "transe.ckpt")) # test the model if test: transe.load_checkpoint(os.path.join(ckpt_path, "transe.ckpt")) tester = Tester(model=transe, data_loader=test_dataloader, use_gpu=True) tester.run_link_prediction(type_constrain=False)
# dataloader for test test_dataloader = TestDataLoader(data_dir, "triple") # define the model transe = TransE(ent_tot=train_dataloader.get_ent_tot(), rel_tot=train_dataloader.get_rel_tot(), dim=embed_dim, p_norm=2, norm_flag=True) # define the loss function model = NegativeSampling(model=transe, loss=MarginLoss(margin=margin), batch_size=train_dataloader.get_batch_size()) # train the model trainer = Trainer(model = model, data_loader = train_dataloader, opt_method = "adam", train_times = train_times, \ alpha = alpha, use_gpu = True, checkpoint_dir=ckpt_path, save_steps=100) tester = Tester(model=transe, data_loader=test_dataloader, use_gpu=True) trainer.run(tester, test_every=100) print("Saving model to {0}...".format(ckpt_path)) transe.save_checkpoint(ckpt_path) # test the model print("Testing...") print("Loading model from {}...".format(ckpt_path)) transe.load_checkpoint(ckpt_path) tester = Tester(model=transe, data_loader=test_dataloader, use_gpu=True) acc, threshlod = tester.run_triple_classification() print("accuracy: ", acc)