train.py

""" Train
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import random
import tqdm

import tensorflow as tf

from utils import start_threads, set_logging_verbosity, MovingAverage, count_number_of_parameters
from data_loader import DATA_PATH, queue_context, tokenize, vectorize
from layers import _phase_train, _phase_infer
from model import Model
from search import reverse_decode, greedy_argmax
from losses import gan_loss

flags = tf.flags

# -- saver and logging options
flags.DEFINE_string("model_dir", "./tmp", (
    "Model directory."))
flags.DEFINE_string("logging_verbosity", "INFO", (
    "Set verbosity to INFO, WARN, DEBUG or ERROR"))

# -- train options
flags.DEFINE_string("corpus_name", "ptb", (
    "Corpus name."))
flags.DEFINE_integer("batch_size", 32, (
    "Batch size for dequeue."))
flags.DEFINE_integer("epoch_size", 10, (
    "Quit after max number of epochs."))
flags.DEFINE_string("seed_text", "how are", (
    "Seed the sampling from the generator with this text."))
flags.DEFINE_string("gan_strategy", "pretrain", (
    "GAN training strategy (pretrain, generator, discriminator, simultaneous, alternating)."))
flags.DEFINE_string("gan_type", "jsd", (
    "GAN type (jsd, emd, ls)."))
flags.DEFINE_float("gan_gd_ratio", 0.5, (
    "Ratio > 0.5 will run generator more, < 0.5 will run discriminator more."))

# -- optimizer options
flags.DEFINE_string("optimizer_type", "adam", (
    "Optimizer type (adam, sgd, rmsprop)."))
flags.DEFINE_float("learning_rate", 1e-4, (
    "Learning rate for optimizer."))
flags.DEFINE_float("learning_rate_decay", 0.8, (
    "Decay the learning rate once criterion is passed."))
flags.DEFINE_float("minimum_learning_rate", 1e-6, (
    "Early stop when lowering than minimum."))
flags.DEFINE_float("max_grads", 5.0, (
    "Max clipping of gradients."))

# -- model options
flags.DEFINE_integer("embedding_dim", 128, (
    "Hidden dimensions for embedding."))
flags.DEFINE_integer("rnn_hidden_dim", 128, (
    "Hidden dimensions for RNN hidden vectors."))
flags.DEFINE_integer("output_hidden_dim", 128, (
    "Hidden dimensions for output hidden vectors before softmax layer."))
flags.DEFINE_float("word_dropout_keep_prob", 0.9, (
    "Dropout keep rate for word embeddings."))
flags.DEFINE_float("recurrent_dropout_keep_prob", 0.6, (
    "Dropout keep rate for recurrent input and output vectors."))
flags.DEFINE_float("output_dropout_keep_prob", 0.5, (
    "Dropout keep rate for output vectors."))

FLAGS = flags.FLAGS
opts = FLAGS.__flags  # dict TODO: make class?

set_logging_verbosity(FLAGS.logging_verbosity)


def _get_n_batches(batch_size, corpus_size):
    return int(corpus_size // batch_size)


def set_initial_ops():
    local_init_op = tf.local_variables_initializer()
    global_init_op = tf.global_variables_initializer()
    init_op = tf.group(local_init_op, global_init_op)
    return init_op


def set_train_op(loss, tvars):
    if FLAGS.optimizer_type == "sgd":
        optimizer = tf.train.GradientDescentOptimizer(learning_rate=FLAGS.learning_rate)
    elif FLAGS.optimizer_type == "rmsprop":
        optimizer = tf.train.RMSPropOptimizer(learning_rate=FLAGS.learning_rate)
    elif FLAGS.optimizer_type == "adam":
        optimizer = tf.train.AdamOptimizer(learning_rate=FLAGS.learning_rate)
    else:
        raise ValueError("Wrong optimizer_type.")

    gradients = optimizer.compute_gradients(loss, var_list=tvars)
    clipped_gradients = [(grad if grad is None else tf.clip_by_norm(grad, FLAGS.max_grads), var)
                         for grad, var in gradients]

    train_op = optimizer.apply_gradients(clipped_gradients)
    return train_op


def get_supervisor(model):
    saver = tf.train.Saver()
    summary_writer = tf.summary.FileWriter(FLAGS.model_dir)

    supervisor = tf.train.Supervisor(
        logdir=FLAGS.model_dir,
        is_chief=True,
        saver=saver,
        init_op=set_initial_ops(),
        summary_op=tf.summary.merge_all(),
        summary_writer=summary_writer,
        save_summaries_secs=100,  # TODO: add as flags
        save_model_secs=1000,
        global_step=model.global_step,
    )

    return supervisor


def get_sess_config():
    # gpu_options = tf.GPUOptions(
    # per_process_gpu_memory_fraction=self.gpu_memory_fraction,
    # allow_growth=True) # seems to be not working

    sess_config = tf.ConfigProto(
        # log_device_placement=True,
        inter_op_parallelism_threads=8,  # TODO: add as flags
        # allow_soft_placement=True,
        # gpu_options=gpu_options)
    )

    return sess_config


def print_loss(sess, loss, moving_average=None):
    l = sess.run(loss)
    if moving_average is None:
        tf.logging.info(" loss: %.4f", l)
    else:
        l_ma = moving_average.next(l)
        tf.logging.info(" loss: %.4f", l_ma)

        # _g, _d = sess.run([g_loss, d_loss])
        # tf.logging.info("g_loss: %.4f, d_loss: %.4f", _g, _d)


# TODO: add to TensorBoard
def print_valid_loss(sess, loss):
    sess.run(_phase_infer)

    total_loss = 0.0
    for _ in range(100):  # TODO: change, use all test data
        l = sess.run(loss)
        total_loss += l

    valid_loss = total_loss / 100.
    tf.logging.info(" valid_loss: %.4f", valid_loss)

    sess.run(_phase_train)


# TODO: configurable seed_text
def print_sample(sess, seed_text, probs, input_ph, word2idx, idx2word):
    # seed_text = "how are you"
    vector = vectorize(seed_text, word2idx)
    out = greedy_argmax(vector[:-1], lambda x: sess.run(probs, {input_ph: [x]}))
    text = reverse_decode(out, idx2word)
    tf.logging.info(" generated text:\n%s", text)


# TODO: learning rate decay
def main():
    corpus = DATA_PATH[FLAGS.corpus_name]
    model = Model(corpus, **opts)

    n_batches = _get_n_batches(FLAGS.batch_size, model.corpus_size)

    # TODO: rename to pretrain
    g_loss = model.g_tensors_pretrain.loss
    g_train_op = set_train_op(g_loss, model.g_tvars)

    g_loss_valid = model.g_tensors_pretrain_valid.loss

    d_logits_real = model.d_tensors_real.prediction_logits
    d_logits_fake = model.d_tensors_fake.prediction_logits

    gan_d_loss, gan_g_loss = gan_loss(
        d_logits_real, d_logits_fake, gan_type=FLAGS.gan_type)

    gan_d_train_op = set_train_op(gan_d_loss, model.d_tvars)
    gan_g_train_op = set_train_op(gan_g_loss, model.g_tvars)

    g_loss_ma = MovingAverage(10)

    sv = get_supervisor(model)
    sess_config = get_sess_config()

    tf.logging.info(" number of parameters %i", count_number_of_parameters())

    with sv.managed_session(config=sess_config) as sess:
        sess.run(_phase_train)

        start_threads(model.enqueue_data, (sess,))
        start_threads(model.enqueue_data_valid, (sess,))

        # TODO: add learning rate decay -> early_stop
        if FLAGS.gan_strategy == "pretrain":
            sv.loop(60, print_loss, (sess, g_loss, g_loss_ma))
            sv.loop(600, print_valid_loss, (sess, g_loss_valid))
            sv.loop(100, print_sample, (sess, FLAGS.seed_text, model.g_tensors_pretrain_valid.flat_logits,
                                        model.input_ph, model.word2idx, model.idx2word))  # TODO: cleanup
        elif FLAGS.gan_strategy in ["generator", "simultaneous", "alternating"]:
            # sv.loop(60, print_loss, (sess, g_loss, g_loss_ma))
            # sv.loop(600, print_valid_loss, (sess, g_loss_valid))
            sv.loop(100, print_sample, (sess, FLAGS.seed_text, model.g_tensors_fake_valid.flat_logits,
                                        model.input_ph, model.word2idx, model.idx2word))

        # make graph read only
        sess.graph.finalize()

        for epoch in range(FLAGS.epoch_size):
            tf.logging.info(" epoch: %i", epoch)

            for _ in tqdm.tqdm(range(n_batches)):
                if sv.should_stop():
                    break

                if FLAGS.gan_strategy == "pretrain":
                    sess.run([g_train_op, model.increment_global_step_op])
                elif FLAGS.gan_strategy == "generator":
                    sess.run([gan_g_train_op, model.increment_global_step_op])
                elif FLAGS.gan_strategy == "discriminator":
                    sess.run([gan_d_train_op, model.increment_global_step_op])
                elif FLAGS.gan_strategy == "simultaneous":
                    sess.run([gan_g_train_op, gan_d_train_op, model.increment_global_step_op])
                elif FLAGS.gan_strategy == "alternating":
                    assert 0. < FLAGS.gan_gd_ratio < 1.0
                    u = random.random()
                    if FLAGS.gan_gd_ratio < u:
                        sess.run([gan_g_train_op, model.increment_global_step_op])
                    elif FLAGS.gan_gd_ratio > u:
                        sess.run([gan_d_train_op, model.increment_global_step_op])
                else:
                    raise ValueError("Wrong gan_strategy.")

                if False:
                    # some criterion
                    sv.stop()

        sv.saver.save(sess, sv.save_path, global_step=sv.global_step)
        tf.logging.info(" training finished")


if __name__ == "__main__":
    main()