launch_simple_tf.py 文件源码

def worker():
  """Worker script that runs on AWS machine. Adds vectors of ones forever,
  prints MB/s."""

  def session_config():
    optimizer_options = tf.OptimizerOptions(opt_level=tf.OptimizerOptions.L0)
    config = tf.ConfigProto(
      graph_options=tf.GraphOptions(optimizer_options=optimizer_options))
    config.operation_timeout_in_ms = 10*1000  # abort after 10 seconds
    return config

  params_size = 250*1000*FLAGS.data_mb # 1MB is 250k floats
  dtype=tf.float32
  val = tf.ones((), dtype=dtype)
  vals = tf.fill([params_size], val)
  params = tf.Variable(vals)
  update = params.assign_add(vals)

  sess = tf.Session(config=session_config())
  sess.run(params.initializer)

  while True:
    start_time = time.perf_counter()
    for i in range(FLAGS.iters_per_step):
      sess.run(update.op)

    elapsed_time = time.perf_counter() - start_time
    rate = float(FLAGS.iters_per_step)*FLAGS.data_mb/elapsed_time
    print('%.2f MB/s'%(rate,))