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')
class Components:
from openke.module.model import TransD
from openke.module.loss import MarginLoss
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 = TransD
loss = MarginLoss(margin=4.0)
strategy = NegativeSampling
trainer = Trainer
tester = Tester
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=25,
neg_rel=0)
# dataloader for test
test_dataloader = TestDataLoader(in_path="./benchmarks/OMKG/",
sampling_mode='link')
# define the model
transe = TransE(ent_tot=train_dataloader.get_ent_tot(),
rel_tot=train_dataloader.get_rel_tot(),
dim=50,
p_norm=1,
norm_flag=True)
model_e = NegativeSampling(model=transe,
loss=MarginLoss(margin=4.0),
batch_size=train_dataloader.get_batch_size())
transr = TransR(ent_tot=train_dataloader.get_ent_tot(),
rel_tot=train_dataloader.get_rel_tot(),
dim_e=50,
dim_r=100,
p_norm=1,
norm_flag=True,
rand_init=False)
model_r = NegativeSampling(model=transr,
loss=MarginLoss(margin=4.0),
batch_size=train_dataloader.get_batch_size())
# pretrain transe
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)
# 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)
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')
neg_ent=64,
neg_rel=0)
# dataloader for test
test_dataloader = TestDataLoader("./benchmarks/CommonGen/", "link")
# define the model
transe = TransE(ent_tot=train_dataloader.get_ent_tot(),
rel_tot=train_dataloader.get_rel_tot(),
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)
# dataloader for test
if test:
test_dataloader = TestDataLoader(bench_path,
"link",
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,
train_dataloader = TrainDataLoader(in_path="./dbpedia50_openKE/kb2E/",
nbatches=100,
threads=8,
bern_flag=1)
# dataloader for test
test_dataloader = TestDataLoader("./dbpedia50_openKE/kb2E/", "link")
# define the model
transe = TransE(ent_tot=train_dataloader.get_ent_tot(),
rel_tot=train_dataloader.get_rel_tot(),
dim=300)
# define the loss function
model = NegativeSampling(model=transe,
loss=MarginLoss(),
batch_size=train_dataloader.get_batch_size())
# train the model
trainer = Trainer(model=model,
data_loader=train_dataloader,
train_times=1000,
alpha=0.01,
use_gpu=True,
opt_method='adagrad')
trainer.run()
transe.save_checkpoint('./checkpoint/transe.ckpt')
embeddings = transe.get_parameters()
directory = './models/dbpedia50/transe300/'
pathlib.Path(directory).mkdir(exist_ok=True, parents=True)
neg_ent=negative_samples,
neg_rel=0)
# 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)
