config = tf.ConfigProto()
config.gpu_options.allow_growth = True
with tf.Session(config=config) as sess:
sess.run(tf.global_variables_initializer())
for epoch in range(pm.epochs):
# train generator
if epoch > 0:
_ = sess.run(model.gen_train_optimizer)
# train critic
disc_loss = 0
for i in range(pm.dis_epochs):
batch = data_loader.next_batch()
disc_loss, _ = sess.run([model.disc_loss, model.disc_train_optimizer], feed_dict={model.real_inputs_discrete: batch})
if epoch % 100 == 0:
samples = []
for i in range(10):
samples.extend(generate_samples(sess, model, inv_charmap))
with open("samples_{}.txt".format(epoch), 'w') as f:
for s in samples:
s = "".join(s)
f.write(s + '\n')
f.flush()
print("MSG : Epoch = [{}/{}], Loss = {}".format(epoch, pm.epochs, disc_loss))
def train(args):
#if args.model_path is None:
# for f in [f for f in os.listdir(args.logdir)]: os.remove(os.path.join(args.logdir, f))
data_loader = Data_loader(batch_size=args.batch_size,
bias_init=args.bias_init)
model = Model(wemb_dim=args.wemb_dim,
hid_dim=args.hid_dim,
seq_len=data_loader.maxlen + 1,
learning_rate=args.learning_rate,
batch_size=args.batch_size,
num_batches=data_loader.num_batches,
num_words=data_loader.num_words,
biivector=data_loader.biivector,
use_gru=args.use_gru)
model.build()
model.loss()
train_op = model.train()
tf.summary.scalar('cross entropy loss', model.cross_entropy_loss_op)
tf.summary.scalar('reg loss', model.reg_loss_op)
tf.summary.scalar('loss', model.loss_op)
merged_op = tf.summary.merge_all()
saver = tf.train.Saver()
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
writer = tf.summary.FileWriter(args.logdir, sess.graph)
sess.run(tf.global_variables_initializer())
if args.model_path is not None:
print('Starting with pretrained model: {}'.format(args.model_path))
saver.restore(sess, args.model_path)
print('Start training')
for ep in xrange(args.epoch):
for step in xrange(data_loader.num_batches):
ctx_batch, cap_batch, mask_batch = data_loader.next_batch()
feed_dict = {
model.ctx_ph:
ctx_batch.reshape(-1, model.ctx_dim[0], model.ctx_dim[1]),
model.cap_ph:
cap_batch,
model.mask_ph:
mask_batch
}
cross_entropy_loss, reg_loss, loss, _, summary = sess.run(
[
model.cross_entropy_loss_op, model.reg_loss_op,
model.loss_op, train_op, merged_op
],
feed_dict=feed_dict)
writer.add_summary(summary, ep * data_loader.num_batches + step)
if step % 100 == 0:
print(
'ep: %2d, step: %4d, xen_loss: %.4f, reg_loss: %.4f, loss: %.4f'
% (ep + 1, step, cross_entropy_loss, reg_loss, loss))
checkpoint_path = os.path.join(args.logdir, 'model.ckpt')
saver.save(sess, checkpoint_path, global_step=ep + 1)
print('Training done.')