コード例 #1
0
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=5000)
    parser.add_argument('--learning_rate', type=float, default=0.003)
    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=10000000)
    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)
    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)
            if (epoch + 1) % 10 == 0:
                kge_model.save_embedding(session=sess)
コード例 #2
0
ファイル: main.py プロジェクト: Governance-com/TransE
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)
コード例 #3
0
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
コード例 #4
0
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
コード例 #5
0
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)
    tf.reset_default_graph()
    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()
        saver.restore(sess, '../checkpoints/gov-g2/GOV-g2130000.ckpt.index')
        # 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('-----Model Loaded-----')
コード例 #6
0
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))
        '''
コード例 #7
0
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)')