mnist.py 文件源码

def main(unused_argv):
  if FLAGS.log_dir is None or FLAGS.log_dir == "":
    raise ValueError("Must specify an explicit `log_dir`")
  if FLAGS.data_dir is None or FLAGS.data_dir == "":
    raise ValueError("Must specify an explicit `data_dir`")

  device, target = device_and_target()
  with tf.device(device):
    images = tf.placeholder(tf.float32, [None, 784], name='image_input')
    labels = tf.placeholder(tf.float32, [None], name='label_input')
    data = read_data_sets(FLAGS.data_dir,
            one_hot=False,
            fake_data=False)
    logits = mnist.inference(images, FLAGS.hidden1, FLAGS.hidden2)
    loss = mnist.loss(logits, labels)
    loss = tf.Print(loss, [loss], message="Loss = ")
    train_op = mnist.training(loss, FLAGS.learning_rate)

  with tf.train.MonitoredTrainingSession(
      master=target,
      is_chief=(FLAGS.task_index == 0),
      checkpoint_dir=FLAGS.log_dir) as sess:
    while not sess.should_stop():
      xs, ys = data.train.next_batch(FLAGS.batch_size, fake_data=False)
      sess.run(train_op, feed_dict={images:xs, labels:ys})