iclr_2017_benchmark.py 文件源码

def main(unused_argv):
  tf.logging.set_verbosity(tf.logging.INFO)
  _logger.info("Tensorflow Version: %s", str(tf.__version__))

  tf_results = test_model(TfModel)
  loom_results = test_model(LoomModel, False)
  loom_results_proper = test_model(LoomModel, True)

  if FLAGS.tree_lstm:
    model_type = "GRU"
  else:
    model_type = "FC"

  _logger.info("====================================================")
  _logger.info("Num epochs: %d; repeats per epoch %d",
               FLAGS.num_epochs, FLAGS.num_repeats)
  _logger.info("Model type: %s, %s", model_type, FLAGS.tree_type)
  _logger.info("Vector size: %d", FLAGS.vector_size)
  _logger.info("Tree size: %d", FLAGS.tree_size)

  print_results(tf_results, "TensorFlow")
  print_results(loom_results, "Loom")
  print_results(loom_results_proper, "Loom with random trees")

  compare_results(tf_results, loom_results, "TensorFlow", "Loom")
  compare_total_speedup(loom_results, tf_results[1])
  _logger.info("Finished benchmarks.")