def main(argv=None):
mnist = input_data.read_data_sets('MNIST_data', one_hot=True)
with tf.device('/gpu:2'):
real_data, z, opt_g, opt_d = build_graph()
summary_op = tf.merge_all_summaries()
saver = tf.train.Saver()
npad = ((0, 0), (2, 2), (2, 2))
with tf.Session() as sess:
sess.run(tf.initialize_all_variables())
summary_writer = tf.train.SummaryWriter(FLAGS.log_dir, sess.graph)
for i in xrange(FLAGS.max_iter_step):
train_img = mnist.train.next_batch(FLAGS.batch_size)[0]
train_img = np.reshape(train_img, (-1, 28, 28))
train_img = np.pad(train_img, pad_width=npad,
mode='constant', constant_values=0)
train_img = np.expand_dims(train_img, -1)
batch_z = np.random.normal(0, 1.0, [FLAGS.batch_size, FLAGS.z_dim]) \
.astype(np.float32)
feed_dict = {real_data: train_img, z: batch_z}
if i % 100 == 99:
run_options = tf.RunOptions(
trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.RunMetadata()
_, merged = sess.run([opt_g, summary_op], feed_dict=feed_dict,
options=run_options, run_metadata=run_metadata)
_, merged = sess.run([opt_g, summary_op], feed_dict=feed_dict,
options=run_options, run_metadata=run_metadata)
summary_writer.add_summary(merged, i)
summary_writer.add_run_metadata(
run_metadata, 'generator_metadata{}'.format(i), i)
_, merged = sess.run([opt_d, summary_op], feed_dict=feed_dict,
options=run_options, run_metadata=run_metadata)
summary_writer.add_summary(merged, i)
summary_writer.add_run_metadata(
run_metadata, 'discriminator_metadata{}'.format(i), i)
else:
sess.run(opt_g, feed_dict=feed_dict)
sess.run(opt_d, feed_dict=feed_dict)
if i % 1000 == 999:
saver.save(sess, os.path.join(
FLAGS.ckpt_dir, "model.ckpt"), global_step=i)
评论列表
文章目录
