text: from dataset import KnowledgeGraph
  from model import TransE
  
  import tensorflow as tf
  import argparse
  
  
  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)
  
  
  if __name__ == '__main__':
      main()
  