def main():
parser = argparse.ArgumentParser(description='TransE')
parser.add_argument('--data_dir', type=str, default='./data/FB15k/')
# parser.add_argument('--embedding_dim', type=int, default=200)
parser.add_argument('--embedding_dim', type=int, default=100)
parser.add_argument('--margin_value', type=float, default=1.0)
parser.add_argument('--score_func', type=str, default='L1')
parser.add_argument('--batch_size', type=int, default=4800)
parser.add_argument('--learning_rate', type=float, default=0.001)
parser.add_argument('--n_generator', type=int, default=6)
parser.add_argument('--n_rank_calculator', type=int, default=6)
parser.add_argument('--ckpt_dir', type=str, default='./ckpt/')
parser.add_argument('--summary_dir', type=str, default='./summary/')
# parser.add_argument('--max_epoch', type=int, default=500)
parser.add_argument('--max_epoch', type=int, default=250)
parser.add_argument('--eval_freq', type=int, default=10)
args = parser.parse_args()
print(args)
kg = KnowledgeGraph(data_dir=args.data_dir)
kge_model = TransH(kg=kg, embedding_dim=args.embedding_dim, margin_value=args.margin_value,
score_func=args.score_func, batch_size=args.batch_size, learning_rate=args.learning_rate,
n_generator=args.n_generator, n_rank_calculator=args.n_rank_calculator)
gpu_config = tf.GPUOptions(allow_growth=True)
sess_config = tf.ConfigProto(gpu_options=gpu_config)
with tf.Session(config=sess_config) as sess:
print('-----Initializing tf graph-----')
tf.global_variables_initializer().run()
print('-----Initialization accomplished-----')
kge_model.check_norm(session=sess)
summary_writer = tf.summary.FileWriter(logdir=args.summary_dir, graph=sess.graph)
for epoch in range(args.max_epoch):
print('=' * 30 + '[EPOCH {}]'.format(epoch) + '=' * 30)
kge_model.launch_training(session=sess, summary_writer=summary_writer)
if (epoch + 1) % args.eval_freq == 0:
kge_model.launch_evaluation(session=sess)
def transe(data_path, embedding_dims, margin_value, score_func, batch_size,
learning_rate, n_generator, n_rank_calculator, max_epoch):
kg = KnowledgeGraph(data_dir=data_path)
kge_model = TransE(kg=kg,
embedding_dim=embedding_dims,
margin_value=margin_value,
score_func=score_func,
batch_size=batch_size,
learning_rate=learning_rate,
n_generator=n_generator,
n_rank_calculator=n_rank_calculator)
gpu_config = tf.GPUOptions(allow_growth=True)
sess_config = tf.ConfigProto(gpu_options=gpu_config)
with tf.Session(config=sess_config) as sess:
print('-----Initializing tf graph-----')
tf.global_variables_initializer().run()
print('-----Initialization accomplished-----')
entity_embedding, relation_embedding = kge_model.check_norm(
session=sess)
summary_writer = tf.summary.FileWriter(logdir='../summary/',
graph=sess.graph)
for epoch in range(max_epoch):
print('=' * 30 + '[EPOCH {}]'.format(epoch) + '=' * 30)
kge_model.launch_training(session=sess,
summary_writer=summary_writer)
if (epoch + 1) % 10 == 0:
kge_model.launch_evaluation(session=sess)
return entity_embedding, relation_embedding
def main():
parser = argparse.ArgumentParser(description='TransE')
parser.add_argument('--data_dir', type=str, default='../data/FB15k/')
parser.add_argument('--embedding_dim', type=int, default=200)
parser.add_argument('--margin_value', type=float, default=1.0)
parser.add_argument('--score_func', type=str, default='L1')
parser.add_argument('--batch_size', type=int, default=4800)
parser.add_argument('--learning_rate', type=float, default=0.001)
parser.add_argument('--n_generator', type=int, default=24)
parser.add_argument('--n_rank_calculator', type=int, default=24)
parser.add_argument('--ckpt_dir', type=str, default='../ckpt/')
parser.add_argument('--summary_dir', type=str, default='../summary/')
parser.add_argument('--max_epoch', type=int, default=500)
parser.add_argument('--eval_freq', type=int, default=10)
args = parser.parse_args()
print(args)
kg = KnowledgeGraph(data_dir=args.data_dir)
kge_model = TransE(kg=kg, embedding_dim=args.embedding_dim, margin_value=args.margin_value,
score_func=args.score_func, batch_size=args.batch_size, learning_rate=args.learning_rate,
n_generator=args.n_generator, n_rank_calculator=args.n_rank_calculator, max_epoch=args.max_epoch,
eval_freq=args.eval_freq)
gpu_config = tf.GPUOptions(allow_growth=True)
sess_config = tf.ConfigProto(gpu_options=gpu_config)
with tf.Session(config=sess_config) as sess:
kge_model.train(sess)
def transe(id,data_path,embedding_dims,margin_value,score_func,batch_size,learning_rate,n_generator,n_rank_calculator,max_epoch,d):
kg = KnowledgeGraph(data_dir=data_path)
content = []
kge_model = TransE(kg=kg, embedding_dim=embedding_dims, margin_value=margin_value,
score_func=score_func, batch_size=batch_size, learning_rate=learning_rate,
n_generator=n_generator, n_rank_calculator=n_rank_calculator,id=id)
gpu_config = tf.GPUOptions(allow_growth=True)
sess_config = tf.ConfigProto(gpu_options=gpu_config)
with tf.Session(config=sess_config) as sess:
# print('-----Initializing tf graph-----')
tf.global_variables_initializer().run()
# print('-----Initialization accomplished-----')
# loss,entity_embedding,relation_embedding = kge_model.check_norm(session=sess)
summary_writer = tf.summary.FileWriter(logdir='../summary/', graph=sess.graph)
for epoch in range(max_epoch):
# print('=' * 30 + '[EPOCH {}]'.format(epoch) + '=' * 30)
# print(loss)
kge_model.launch_training(session=sess, summary_writer=summary_writer)
if (epoch + 1) % 50 == 0:
kge_model.launch_evaluation(session=sess)
loss,entity_embedding,relation_embedding = kge_model.check_norm(session=sess)
content.append(loss)
content.append(entity_embedding)
content.append(relation_embedding)
# print(relation_embedding.shape)
d[id] = content
print('FB15k-{} loss:{}'.format(id,d[id][0]))
# print(type(d))
return entity_embedding,relation_embedding
def main():
parser = argparse.ArgumentParser(description='TransE')
parser.add_argument('--data_dir', type=str, default='../data/FB15k/')
parser.add_argument('--embedding_dim', type=int, default=200)
parser.add_argument('--margin_value', type=float, default=1.0)
parser.add_argument('--score_func', type=str, default='L1')
parser.add_argument('--batch_size', type=int, default=4800)
parser.add_argument('--learning_rate', type=float, default=0.001)
parser.add_argument('--n_generator', type=int, default=24)
parser.add_argument('--n_rank_calculator', type=int, default=24)
parser.add_argument('--ckpt_dir', type=str, default='../ckpt/')
parser.add_argument('--model_name', type=str)
parser.add_argument('--summary_dir', type=str, default='../summary/')
parser.add_argument('--max_epoch', type=int, default=500)
parser.add_argument('--eval_freq', type=int, default=10)
args = parser.parse_args()
print(args)
kg = KnowledgeGraph(data_dir=args.data_dir)
kge_model = TransE(kg=kg,
embedding_dim=args.embedding_dim,
margin_value=args.margin_value,
score_func=args.score_func,
batch_size=args.batch_size,
learning_rate=args.learning_rate,
n_generator=args.n_generator,
n_rank_calculator=args.n_rank_calculator,
model_name=args.model_name,
ckpt_dir=args.ckpt_dir)
gpu_config = tf.GPUOptions(allow_growth=False)
sess_config = tf.ConfigProto(gpu_options=gpu_config)
with tf.Session(config=sess_config) as sess:
print('-----Initializing tf graph-----')
tf.global_variables_initializer().run()
print('-----Initialization accomplished-----')
kge_model.check_norm(session=sess)
summary_writer = tf.summary.FileWriter(logdir=args.summary_dir,
graph=sess.graph)
saver = tf.train.Saver(tf.global_variables(), max_to_keep=500)
for epoch in range(args.max_epoch):
print('=' * 30 + '[EPOCH {}]'.format(epoch) + '=' * 30)
kge_model.launch_training(session=sess,
summary_writer=summary_writer)
if (epoch + 1) % args.eval_freq == 0:
kge_model.launch_evaluation(session=sess, saver=saver)
print('-----Save checkpoint-----')
step_str = str(kge_model.global_step.eval(session=sess))
save_path = args.ckpt_dir + '/' + args.model_name + step_str + '.ckpt'
saver_path = saver.save(sess, save_path)
tf.saved_model.simple_save(
sess,
args.ckpt_dir + '/model-' + step_str,
inputs={'triple': kge_model.eval_triple},
outputs={
'entity-embedding': kge_model.entity_embedding,
'relation-embedding': kge_model.relation_embedding
})
print("Model saved in path: %s" % saver_path)
def accuracy(params, session=sess):
kg = KnowledgeGraph(data_dir=args.data_dir)
kge_model = RelationPredict(kg=kg,
batch_size=params['batch_size'],
learning_rate=params['learning_rate'])
tf.global_variables_initializer().run()
summary_writer = tf.summary.FileWriter(logdir=args.summary_dir,
graph=session.graph)
for epoch in range(args.epoch_num):
kge_model.launch_training(session=session,
summary_writer=summary_writer)
return -kge_model.epoch_accuracy
def main():
parser = argparse.ArgumentParser(description='TransE')
parser.add_argument('--data_dir', type=str, default='../data/')
parser.add_argument('--embedding_dim', type=int, default=200)
parser.add_argument('--margin_value', type=float, default=1.0)
parser.add_argument('--score_func', type=str, default='L1')
parser.add_argument('--batch_size', type=int, default=4800)
parser.add_argument('--learning_rate', type=float, default=0.001)
parser.add_argument('--n_generator', type=int, default=24)
parser.add_argument('--n_rank_calculator', type=int, default=24)
parser.add_argument('--ckpt_dir', type=str, default='../ckpt/')
parser.add_argument('--summary_dir', type=str, default='../summary/')
parser.add_argument('--max_epoch', type=int, default=500)
parser.add_argument('--eval_freq', type=int, default=10)
args = parser.parse_args()
print(args)
'''
传递参数
Namespace(
batch_size=4800,
ckpt_dir='../ckpt/',
data_dir='../data/FB15k/',
embedding_dim=200,
eval_freq=10, learning_rate=0.001,
margin_value=1.0,
max_epoch=500,
n_generator=24,
n_rank_calculator=24,
score_func='L1',
//summary_dir='../summary/')
'''
kg = KnowledgeGraph(data_dir=args.data_dir)#create graph
kge_model = TransE(
kg=kg,
embedding_dim=args.embedding_dim,
margin_value=args.margin_value,
score_func=args.score_func,
batch_size=args.batch_size,
learning_rate=args.learning_rate,
n_generator=args.n_generator,
n_rank_calculator=args.n_rank_calculator
)#init embd...
gpu_config = tf.GPUOptions(allow_growth=True)
'''
tf.GPUOptions:可以作为设置tf.ConfigProto时的一个参数选项，一般用于限制GPU资源的使用
allow_growth=True:动态申请现显存
'''
sess_config = tf.ConfigProto(gpu_options=gpu_config)
'''
tf.ConfigProto:
创建session的时候，用来对session进行参数配置
'''
with tf.Session(config = sess_config) as sess:
'''
Session 是 Tensorflow 为了控制,和输出文件的执行的语句. 运行 session.run() 可以获得你要得知的运算结果
'''
print('-----Initializing tf graph-----')
tf.global_variables_initializer().run()#就是 run了 所有global Variable 的 assign op，这就是初始化参数的本来面目。
print('-----Initialization accomplished-----')
kge_model.check_norm(session=sess)
summary_writer = tf.summary.FileWriter(logdir=args.summary_dir, graph=sess.graph)
# print(type(sess.graph))
# print(type(summary_writer))
'''
relation) in self.kg.golden_triple_pool:
continue
else:
tail_rank_filter += 1
out_queue.put((head_rank_raw, tail_rank_raw, head_rank_filter,
tail_rank_filter))
in_queue.task_done()
def test(self, in_queue, out_queue):
print("happy!")
if __name__ == '__main__':
freeze_support()
args = None
kg = KnowledgeGraph('../data/FB15k')
# transe = TransE(kg.n_entity, kg.n_relation, dim =200)
demo = TransEDemo(kg, args)
training_range = tqdm(range(1000)) # 进度条
for epoch in training_range:
start = time.time()
epoch_loss = demo.launch_training(epoch + 1)
#training_range.set_description('Epoch {}, epoch loss: {:.4f}, cost time: {:.4f}s'.format(epoch, epoch_loss,
print('Epoch {}, epoch loss: {:.4f}, cost time: {:.4f}s'.format(
epoch, epoch_loss,
time.time() - start))
demo.model.save_checkpoint('checkpoint/model_params.pkl')
demo.model.load_checkpoint('checkpoint/model_params.pkl')
demo.model = demo.model.cpu()
demo.launch_evaluation()
batch_loss / len(curr_node_batch), batch_accuracy),
end='\r')
self.epoch_accuracy = epoch_accuracy / batch_num
print()
print('epoch loss: {:.3f}'.format(epoch_loss))
print('epoch accuracy: {}'.format(self.epoch_accuracy))
print('cost time: {:.3f}s'.format(timeit.default_timer() - start))
print('-----------------Finish training-----------------')
def check(self, session):
print('-----------Check----------')
rela = self.relation_matrix.eval(session=session)
bias = self.bias.eval(session=session)
print('relation matrix:')
print(rela)
print('bias:')
print(bias)
print('--------Check Finished----')
if __name__ == '__main__':
kg = KnowledgeGraph(data_dir='../data/')
kg.training_num = 4
model = RelationPredict(kg=kg, batch_size=4, learning_rate=0.01)
sess = tf.Session()
sess.run(tf.global_variables_initializer())
summary_writer = tf.summary.FileWriter(logdir='../summ', graph=sess.graph)
model.launch_training(session=sess, summary_writer=summary_writer)
model.check(session=sess)
sess.close()
import math
from dataset import KnowledgeGraph
pp = pprint.PrettyPrinter(indent=4)
# init_op = tf.initialize_all_variables()
# init_op_t = tf.initialize_all_tables()
# saver = tf.train.Saver()
ent_data = open('/Users/mo/Downloads/Entity.json').read()
gov_ents = json.loads(ent_data)
l = gov_ents.get('records')
l.insert(0, {"id": "63fbed08-b4b5-4881-9636-e7c18d1fccda"})
kg = KnowledgeGraph(data_dir='./data/GOV/')
builder = tf.saved_model.builder.SavedModelBuilder('./GOV-g2/model-98800/')
with tf.Session() as sess:
print('------started session------')
# tf.saved_model.loader.load(sess, ['serve'], )
builder.add_meta_graph_and_variables(
sess,
['serve'],
# signature_def_map=foo_signatures,
# assets_collection=foo_assets,
strip_default_attrs=True)
print('------model loaded------')
def main():
parser = argparse.ArgumentParser(description='TransE')
parser.add_argument('--mode', type=str, default='eval')
parser.add_argument('--data_dir', type=str, default='../data/FB15k/')
parser.add_argument('--embedding_dim', type=int, default=200)
parser.add_argument('--margin_value', type=float, default=1.0)
parser.add_argument('--score_func', type=str, default='L1')
parser.add_argument('--batch_size', type=int, default=4800)
parser.add_argument('--learning_rate', type=float, default=0.001)
parser.add_argument('--n_generator', type=int, default=24)
parser.add_argument('--n_rank_calculator', type=int, default=24)
parser.add_argument('--hit_at_n', type=int, default=10)
parser.add_argument('--ckpt_dir', type=str, default='../ckpt/')
parser.add_argument('--summary_dir', type=str, default='../summary/')
parser.add_argument('--max_epoch', type=int, default=500)
parser.add_argument('--eval_freq', type=int, default=10)
args = parser.parse_args()
print(args)
kg = KnowledgeGraph(data_dir=args.data_dir)
kge_model = TransE(kg=kg,
model_path=args.ckpt_dir,
embedding_dim=args.embedding_dim,
margin_value=args.margin_value,
score_func=args.score_func,
batch_size=args.batch_size,
learning_rate=args.learning_rate,
n_generator=args.n_generator,
n_rank_calculator=args.n_rank_calculator,
hit_at_n=args.hit_at_n)
gpu_config = tf.GPUOptions(allow_growth=True)
sess_config = tf.ConfigProto(gpu_options=gpu_config)
if args.mode == 'test':
saver = tf.train.Saver()
else:
saver = tf.train.Saver(tf.global_variables())
with tf.Session(config=sess_config) as sess:
if args.mode == 'eval':
print('-----Loading from checkpoints-----')
ckpt_file = tf.train.latest_checkpoint(args.ckpt_dir)
saver.restore(sess, ckpt_file)
kge_model.launch_evaluation(session=sess)
elif args.mode == 'train':
print('-----Initializing tf graph-----')
tf.global_variables_initializer().run()
print('-----Initialization accomplished-----')
kge_model.check_norm(session=sess)
summary_writer = tf.summary.FileWriter(logdir=args.summary_dir,
graph=sess.graph)
for epoch in range(args.max_epoch):
print('=' * 30 + '[EPOCH {}]'.format(epoch) + '=' * 30)
kge_model.launch_training(epoch=epoch,
session=sess,
summary_writer=summary_writer,
saver=saver)
# if (epoch + 1) % args.eval_freq == 0:
# kge_model.launch_evaluation(session=sess)
elif args.mode == 'predict':
print('-----Loading from checkpoints-----')
ckpt_file = tf.train.latest_checkpoint(args.ckpt_dir)
saver.restore(sess, ckpt_file)
kge_model.launch_prediction(session=sess)
else:
print('Wrong mode!! (mode = train|test|predict)')
import numpy as np
from dataset import KnowledgeGraph
from predict_model import TransE
import tensorflow as tf
from xmlrpc.server import SimpleXMLRPCServer
kg = KnowledgeGraph(data_dir='../../Data/TrainingData/')
ent_dic = kg.entity_dict
rel_dic = kg.relation_dict
model = TransE(kg=kg, score_func='L1')
type_dic = dict()
with open('../../Data/TrainingData/type.txt', 'r', encoding='utf-8') as f:
for i in f.readlines():
s = i.strip().split('\t')
type_dic[s[0]] = s[1]
def predict_drug(t):
r_id = int(rel_dic.get('effect'))
t_id = int(ent_dic.get(t))
top_k_id = model.predict_head(r_id, t_id)
top_k = list()
new_ent_dic = dict(zip(ent_dic.values(), ent_dic.keys()))
sess = tf.Session()
with sess.as_default():
for i in top_k_id.eval():
top_k.append(new_ent_dic.get(i))
print()
result = list()
count = 0
for name in top_k: