def main():
mnist = input_data.read_data_sets('MNIST_data', one_hot=True)
with tf.device('/gpu:1'):
g_loss_sum, d_loss_sum, img_sum, opt_g, opt_d, z, real_data = build_graph()
summary_g = tf.merge_summary([g_loss_sum, img_sum])
summary_d = tf.merge_summary([d_loss_sum, img_sum])
saver = tf.train.Saver()
npad = ((0, 0), (2, 2), (2, 2))
with tf.Session(config=tf.ConfigProto(
allow_soft_placement=True)) 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_data = mnist.train.next_batch(FLAGS.batch_size)
train_img = np.reshape(train_data[0], (-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.uniform(-1, 1, [FLAGS.batch_size, FLAGS.z_dim]) \
.astype(np.float32)
feed_dict = {real_data[0]: train_img, z: batch_z, real_data[1]:train_data[1]}
if i % 100 == 99:
run_options = tf.RunOptions(
trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.RunMetadata()
_, merged = sess.run([opt_g, summary_g], 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_g, summary_g], feed_dict=feed_dict,
options=run_options, run_metadata=run_metadata)
summary_writer.add_summary(merged, i)
summary_writer.add_run_metadata(
run_metadata, 'second_generator_metadata {}'.format(i), i)
_, merged = sess.run([opt_d, summary_d], 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_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)
评论列表
文章目录
