def main():
# Define some hyper-parameters for training
emb_dim = 100
lr = 0.001 #0.0004
margin = 1
n_epochs = 30 #10
batch_size = 5000 #10000 #32768
# Load dataset
#kg_train, kg_val, kg_test = load_fb15k()
#kinship_df = pd.read_csv('data/kinship.txt', delimiter='\t', header=None, names=['from', 'rel', 'to'])
#kinship_kg = KnowledgeGraph(df=kinship_df)
#kg_train, kg_test = kinship_kg.split_kg(share=0.9)
data_name = 'FB15k-237'
df1 = pd.read_csv('../data/%s/divided/train1.csv' % data_name,
delimiter='\t')
#df1 = df1.rename(columns={'head': 'from', 'rel': 'rel', 'tail': 'to'})
kg = KnowledgeGraph(df1)
#kg_train, kg_test = kg.split_kg(share=0.85)
kg_train, kg_test = kg.split_kg(size=(0.85, ))
print('n_ent: ', kg_train.n_ent)
print('n_rel: ', kg_train.n_rel)
print('n_facts: train: %s, test: %s' % (kg_train.n_facts, kg_test.n_facts))
#print('n_facts: total: %s, train: %s, test: %s' %(kinship_kg.n_facts, kg_train.n_facts, kg_test.n_facts))
# Define the model and criterion
#model = RESCALModel(emb_dim, kg_train.n_ent, kg_train.n_rel)
#print('RESCALModel')
model = Rescal(kg_train.n_ent, kg_train.n_rel, emb_dim)
criterion = MarginLoss(margin)
optimizer = Adam(model.parameters(), lr=lr, weight_decay=1e-5)
trainer = Trainer(model,
criterion,
kg_train,
n_epochs,
batch_size,
optimizer=optimizer,
sampling_type='unif',
use_cuda=None)
for _ in range(1):
trainer.run()
evaluator = LinkPredictionEvaluator(model, kg_test)
evaluator.evaluate(200, 10)
evaluator.print_results(k=[1, 3, 10])
def main_quantified_TransE():
# Define some hyper-parameters for training
emb_dim = 100
lr = 0.0004
margin = 0.5
n_epochs = 1000
batch_size = 2097152
# Load dataset
data_path = "/tmp/pycharm_project_583/data/uncmtrd/agg6_202005_ALL_tv.csv"
kg_train, kg_val, kg_test = load_custom_qr(data_path=data_path)
model = TransEQuantifiedRelations(
emb_dim, kg_train.n_ent, kg_train.n_rel, dissimilarity_type="L2"
)
print_model(model) # check we only have two embedding layers - one for entity, the other for relations
dataset_name = data_path.split('/')[-1].replace('.csv', '')
curr_time = datetime.now().strftime('%Y%m%d%H%M%S')
model_prefix = os.path.join('./pretrained', f'{dataset_name}_emb{emb_dim}_lr{lr}_mgn{margin}_epch{n_epochs}_bsize{batch_size}_t{curr_time}')
criterion = MarginLoss(margin)
optimizer = Adam(model.parameters(), lr=lr, weight_decay=1e-5)
trainer = Trainer(
model,
criterion,
kg_train,
n_epochs,
batch_size,
optimizer=optimizer,
sampling_type="bern",
use_cuda=None,
)
trainer.run(kg_test=kg_test, model_prefix=model_prefix)
def main():
# Define some hyper-parameters for training
emb_dim = 100
lr = 0.0004
margin = 0.5
n_epochs = 1000
batch_size = 32768
# Load dataset
kg_train, kg_val, kg_test = load_fb15k()
print(kg_train)
# Define the model and criterion
model = TransEModel(emb_dim,
kg_train.n_ent,
kg_train.n_rel,
dissimilarity_type="L2")
criterion = MarginLoss(margin)
optimizer = Adam(model.parameters(), lr=lr, weight_decay=1e-5)
trainer = Trainer(
model,
criterion,
kg_train,
n_epochs,
batch_size,
optimizer=optimizer,
sampling_type="bern",
use_cuda=None,
)
trainer.run()
evaluator = LinkPredictionEvaluator(model, kg_test)
evaluator.evaluate(200, 10)
evaluator.print_results()
def main():
# Define some hyper-parameters for training
global optimizer
benchmarks = 'GeoDBpedia21'
model_name = 'TransR_GDR'
opt_method = 'Adam' # "Adagrad" "Adadelta" "Adam" "SGD"
GDR = True # 是否引入坐标信息
emb_dim = 100 # TransE model
ent_dim = emb_dim
rel_dim = emb_dim
lr = 0.001
margin = 0.5
n_epochs = 20000
train_b_size = 256 # 训练时batch size
eval_b_size = 64 # 测评valid test 时batch size
# save_time_freq = 5
# require_improvement = save_time_freq*5
validation_freq = 10 # 多少轮进行在验证集进行一次测试 同时保存最佳模型
require_improvement = validation_freq * 3 # 验证集top_k超过多少epoch没下降，结束训练
model_save_path = './checkpoint/' + benchmarks + '_' + model_name + '_' + opt_method + '.ckpt' # 保存最佳hits k (ent)模型
device = 'cuda:0' if cuda.is_available() else 'cpu'
# Load dataset
module = getattr(import_module('torchkge.models'), model_name + 'Model')
load_data = getattr(import_module('torchkge.utils.datasets'), 'load_' + benchmarks)
print('Loading data...')
kg_train, kg_val, kg_test = load_data(GDR=GDR)
print(f'Train set: {kg_train.n_ent} entities, {kg_train.n_rel} relations, {kg_train.n_facts} triplets.')
print(f'Valid set: {kg_val.n_facts} triplets, Test set: {kg_test.n_facts} triplets.')
# Define the model and criterion
print('Loading model...')
if 'TransE' in model_name:
model = module(emb_dim, kg_train.n_ent, kg_train.n_rel, dissimilarity_type='L2')
else:
model = module(ent_dim, rel_dim, kg_train.n_ent, kg_train.n_rel)
criterion = MarginLoss(margin)
# Move everything to CUDA if available
if device == 'cuda:0':
cuda.empty_cache()
model.to(device)
criterion.to(device)
dataloader = DataLoader(kg_train, batch_size=train_b_size, use_cuda='all')
else:
dataloader = DataLoader(kg_train, batch_size=train_b_size, use_cuda=None)
# Define the torch optimizer to be used
optimizer = optimizer(model, opt_method=opt_method, lr=lr)
# optimizer = Adam(model.parameters(), lr=lr, weight_decay=1e-5)
sampler = BernoulliNegativeSampler(kg_train)
start_epoch = 1
best_score = float('-inf')
if os.path.exists(model_save_path): # 存在则加载模型 并继续训练
start_epoch, best_score = load_ckpt(model_save_path, model, optimizer)
print(f'loading ckpt sucessful, start on epoch {start_epoch}...')
print(model)
print('lr: {}, margin: {}, dim {}, total epoch: {}, device: {}, batch size: {}, optim: {}, GDR: {}' \
.format(lr, margin, emb_dim, n_epochs, device, train_b_size, opt_method, GDR))
print('Training...')
last_improve = start_epoch # 记录上次验证集loss下降的epoch数
start = time.time()
# last_improve = start
# save_time = start
for epoch in range(start_epoch, n_epochs + 1):
# model.normalize_parameters()
running_loss = 0.0
model.train()
for i, batch in enumerate(dataloader):
if GDR:
h, t, r, point = batch[0], batch[1], batch[2], batch[3]
n_h, n_t = sampler.corrupt_batch(h, t, r) # 1:1 negative sampling
n_point = id2point(n_h, n_t, kg_train.id2point)
optimizer.zero_grad()
# forward + backward + optimize
pos, neg = model(h, t, n_h, n_t, r)
loss = criterion(pos, neg, point, n_point)
else:
h, t, r = batch[0], batch[1], batch[2]
n_h, n_t = sampler.corrupt_batch(h, t, r)
optimizer.zero_grad()
pos, neg = model(h, t, n_h, n_t, r)
loss = criterion(pos, neg)
loss.backward()
optimizer.step()
running_loss += loss.item()
model.normalize_parameters()
# print('\rEpoch [{:>4}/{:>4}] | mean loss: {:>8.3f}, time: {}'.format(epoch, n_epochs, running_loss / len(dataloader), time_since(start)), end='', flush=True)
# # test
if epoch % validation_freq == 0:
create_dir_not_exists('./checkpoint')
model.eval()
evaluator = LinkPredictionEvaluator(model, kg_val)
evaluator.evaluate(b_size=eval_b_size, verbose=False)
_, hit_at_k = evaluator.hit_at_k(10) # val filter hit_k
print('Epoch [{:>5}/{:>5}] '.format(epoch, n_epochs), end='')
if hit_at_k > best_score:
save_ckpt(model, optimizer, epoch, best_score, model_save_path)
best_score = hit_at_k
improve = '*' # 在有提升的结果后面加上*标注
last_improve = epoch # 验证集hit_k增大即认为有提升
else:
improve = ''
msg = '| Train loss: {:>8.3f}, Val Hit@10: {:>5.2%}, Time {} {}'
print(msg.format(running_loss / len(dataloader), hit_at_k, time_since(start), improve))
# model.normalize_parameters()
if epoch - last_improve > require_improvement:
# 验证集top_k超过一定epoch没增加，结束训练
print("\nNo optimization for a long time, auto-stopping...")
break
# # test
# if (time.time() - save_time)/60 > save_time_freq:
# create_dir_not_exists('./checkpoint')
# model.eval()
# evaluator = LinkPredictionEvaluator(model, kg_val)
# evaluator.evaluate(b_size=eval_b_size, verbose=False)
# _, hit_at_k = evaluator.hit_at_k(10) # val filter hit_k
# if hit_at_k > best_score:
# save_ckpt(model, optimizer, epoch, best_score, model_save_path)
# best_score = hit_at_k
# improve = '*' # 在有提升的结果后面加上*标注
# last_improve = time.time() # 验证集hit_k增大即认为有提升
# else:
# improve = ''
# save_time = time.time()
# msg = ', Val Hit@10: {:>5.2%} {}'
# print(msg.format(hit_at_k, improve))
# model.normalize_parameters()
# if (time.time() - last_improve)/60 > require_improvement:
# # 验证集top_k超过一定epoch没增加，结束训练
# print("\nNo optimization for a long time, auto-stopping...")
# break
print('\nTraining done, start evaluate on test data...')
print('model name: {}, lr: {}, dim {}, device: {}, eval batch size: {}, optim: {}, GDR: {}' \
.format(model_name, lr, emb_dim, device, eval_b_size, opt_method, GDR))
# Testing the best checkpoint on test dataset
load_ckpt(model_save_path, model, optimizer)
model.eval()
lp_evaluator = LinkPredictionEvaluator(model, kg_test)
lp_evaluator.evaluate(eval_b_size, verbose=False)
lp_evaluator.print_results()
rp_evaluator = RelationPredictionEvaluator(model, kg_test)
rp_evaluator.evaluate(eval_b_size, verbose=False)
rp_evaluator.print_results()
print(f'Total time cost: {time_since(start)}')
# Define some hyper-parameters for training
emb_dim = 250
lr = 0.0004
n_epochs = 1000
b_size = 32768
margin = 0.5
# In[18]:
# Define the model and criterion
model = TransEModel(emb_dim,
kg_train.n_ent,
kg_train.n_rel,
dissimilarity_type='L2')
criterion = MarginLoss(margin)
# In[43]:
# Move everything to CUDA if available
if cuda.is_available():
cuda.empty_cache()
model.cuda()
criterion.cuda()
# In[20]:
# Define the torch optimizer to be used
optimizer = Adam(model.parameters(), lr=lr, weight_decay=1e-5)
sampler = BernoulliNegativeSampler(kg_train)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if device.type == 'cuda':
print('gpu is available')
torch.cuda.empty_cache()
if torch.cuda.device_count() > 1:
print('multiple gpus are available')
if args.gpu is not None:
model = DataParallel(model, device_ids=args.gpu)
else:
model = DataParallel(model)
checkpoint_manager = CheckpointManager(restore_dir)
ckpt = checkpoint_manager.load_checkpoint(f'best_{args.model}.tar')
model.load_state_dict(ckpt['model_state_dict'])
criterion = MarginLoss(margin)
model.to(device)
criterion.to(device)
sampler = BernoulliNegativeSampler(kg_test)
test_dl = DataLoader(kg_test, batch_size=args.batch_size)
model.eval()
test_loss = 0
for step, batch in tqdm(enumerate(test_dl),
desc='steps',
total=len(test_dl)):
h, t, r = map(lambda elm: elm.to(device), batch)
n_h, n_t = sampler.corrupt_batch(h, t, r)
with torch.no_grad():
assert args.model in ['TransE', 'TransR', 'DistMult'
], "Invalid Knowledge Graph Embedding Model"
if args.model == 'TransE':
model = torchkge.models.TransEModel(args.ent_dim,
kg_train.n_ent,
kg_train.n_rel,
dissimilarity_type='L2')
elif args.model == 'DistMult':
model = torchkge.models.DistMultModel(args.ent_dim, kg_train.n_ent,
kg_train.n_rel)
elif args.model == 'TransR':
model = torchkge.models.TransRModel(args.ent_dim, args.rel_dim,
kg_train.n_ent, kg_train.n_rel)
criterion = MarginLoss(args.margin)
if torch.cuda.is_available():
torch.cuda.empty_cache()
model.cuda()
criterion.cuda()
writer = SummaryWriter(save_dir / f'runs_{args.model}')
checkpoint_manager = CheckpointManager(save_dir)
summary_manager = SummaryManager(save_dir)
summary_manager.update(experiment_summary)
optimizer = torch.optim.Adam(model.parameters(),
lr=args.learning_rate,
weight_decay=1e-5)
sampler = BernoulliNegativeSampler(kg_train)