def main(_):
_model = wgan_gp()
_gpu = tf.GPUOptions(allow_growth=True)
_saver = tf.train.Saver(pad_step_number=True)
with tf.Session(config=tf.ConfigProto(gpu_options=_gpu)) as sess:
_writer = tf.summary.FileWriter(FLAGS.log_path, sess.graph)
tf.global_variables_initializer().run()
ckpt = tf.train.get_checkpoint_state(FLAGS.log_path)
if ckpt and ckpt.model_checkpoint_path:
_saver.restore(sess, FLAGS.log_path)
_step = global_step.eval()
while True:
if _step >= FLAGS.steps:
break
d_loss, g_loss, fit_summary = _model.fit(sess, 100)
_step = _step + 100
_writer.add_summary(fit_summary, _step)
print("Train [%d\%d] g_loss [%3f] d_loss [%3f]" %
(_step, FLAGS.steps, g_loss, d_loss))
images = _model.gen(sess, 100)
imsave_(FLAGS.log_path + 'train_{}.png'.format(_step),
imcombind_(images))
_saver.save(sess, FLAGS.log_path) if _step % 5000 == 0 else None
max_outputs=4)
])
# train
with tf.Session() as sess:
tf.global_variables_initializer().run()
_writer = tf.summary.FileWriter(log_path, sess.graph)
for it in range(n_iters):
x, _ = mnist.train.next_batch(batch_size, shuffle=True)
noise = np.random.normal(size=[batch_size, z_dim])
sess.run(d_opt, {x_real: x, z: noise})
sess.run(g_opt, {x_real: x, z: noise})
if it % to_display == 0:
res = sess.run([
g_loss, d_loss,
tf.reduce_mean(real_score),
tf.reduce_mean(fake_score), x_fake, summery
], {
x_real: x,
z: noise
})
_writer.add_summary(res[5], it)
im = imcombind_(np.reshape(res[4], [batch_size, 28, 28, 1]))
imsave_(log_path + '%d.png' % it, im)
print(
'%4d/%4d g_loss: %.3f d_loss %.3f real_score: %.3f fake_score: %.3f '
% (it, n_iters, res[0], res[1], res[2], res[3]))