示例#1
0
文件: train.py 项目: linzeqipku/HSP
def main(args):
    # log level
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    tf.logging.set_verbosity(tf.logging.INFO)

    # model & params & dirs
    model_cls = decomp_models.get_decomp_model(args.model)
    params = parse_params(args, default_parameters(), model_cls.get_parameters())
    best_model_dir = prepare_dir(params)

    # tf.set_random_seed(params.seed)

    # define evaluation args & op
    ARGS = namedtuple('Args', ['input', 'output', 'vocab', 'models', 'model', 'parameters', 'log'])
    trans_file = os.path.join(params.output, os.path.basename(params.validation) + '.trans')
    eval_args = ARGS(
        input=params.validation,
        output=trans_file,
        vocab=[params.vocab[0], params.vocab[1]],
        models=[params.output, ],
        model=args.model,
        parameters=params.dev_params,
        log=False
    )

    def eval_op():
        tf.logging.info("Evaluate model on dev set...")
        inference.main(eval_args, verbose=False)
        return eval_args.output, evaluate(pred_file=eval_args.output, ref_file=args.references)

    # Build Graph
    with tf.Graph().as_default():
        # Build input queue
        batcher = Batcher(params, 'train')

        # Build model
        initializer = get_initializer(params)
        model = model_cls(params, args.model, initializer=initializer)

        # Create global step
        global_step = tf.train.get_or_create_global_step()

        # Print trainable parameter and their shape
        all_weights = {v.name: v for v in tf.trainable_variables()}
        for v_name in sorted(list(all_weights)):
            v = all_weights[v_name]
            tf.logging.info("%s\tshape    %s", v.name[:-2].ljust(80),
                            str(v.shape).ljust(20))
        # Print parameter number
        total_size = sum([np.prod(v.shape.as_list()) for v in tf.trainable_variables()])
        tf.logging.info("Total trainable variables size: %d", total_size)

        # lr decay
        learning_rate = get_learning_rate_decay(params.learning_rate,
                                                global_step, params)
        if params.learning_rate_minimum:
            lr_min = float(params.learning_rate_minimum)
            learning_rate = tf.maximum(learning_rate, tf.to_float(lr_min))

        learning_rate = tf.convert_to_tensor(learning_rate, dtype=tf.float32)
        tf.summary.scalar("learning_rate", learning_rate)

        # Create optimizer
        if params.optimizer == "Adam":
            opt = tf.train.AdamOptimizer(learning_rate,
                                         beta1=params.adam_beta1,
                                         beta2=params.adam_beta2,
                                         epsilon=params.adam_epsilon)
        elif params.optimizer == "LazyAdam":
            opt = tf.contrib.opt.LazyAdamOptimizer(learning_rate,
                                                   beta1=params.adam_beta1,
                                                   beta2=params.adam_beta2,
                                                   epsilon=params.adam_epsilon)
        else:
            raise RuntimeError("Optimizer %s not supported" % params.optimizer)

        loss, ops = optimize.create_train_op(model.loss, opt, global_step, params)

        init_op = init_variables()

        restore_op = restore_variables(args.output)

        config = session_config(params)

        best_metric = 0
        saver = tf.train.Saver(tf.trainable_variables(), max_to_keep=params.keep_checkpoint_max)
        best_saver = tf.train.Saver(tf.trainable_variables(), max_to_keep=params.keep_checkpoint_max)
        with tf.Session(config=config) as sess:
            sess.run(init_op)
            sess.run(restore_op)
            t0 = time.time()
            while True:
                batch = batcher.next_batch()
                _, step = sess.run([ops, global_step], model.make_feed_dict(batch))

                if step > params.train_steps:  # terminate
                    tf.logging.info('Train done.')
                    break

                if params.save_checkpoint_steps and step % params.save_checkpoint_steps == 0:  # save checkpoint
                    model.save(saver, sess, params.output, model_prefix='{}-{}'.format(args.model, step))

                if step % params.print_steps == 0:  # print message
                    t1 = time.time()
                    loss = sess.run(model.loss, model.make_feed_dict(batch))
                    tf.logging.info('seconds for training step: %.3f', t1 - t0)
                    t0 = time.time()
                    tf.logging.info('Step {}, Loss {}'.format(step, loss))

                if step % params.eval_steps == 0 and step > params.eval_steps_begin:  # evaluation
                    trans_output, (total, comp_acc, metrics, em) = eval_op()
                    bleu, rouge = metrics['Bleu-4'], metrics['Rouge-L']
                    f1 = 2 * bleu * rouge / (bleu + rouge + 1e-6)
                    metric = f1 if args.metric == 'f1' else em
                    if metric > best_metric:
                        best_metric = metric
                        tf.logging.info("Step {} -> best model acc: {:.3f}, bleu-4: {:.3f}, rouge-L: {:.3f}.".format(
                            step, comp_acc, bleu, rouge))
                        tf.logging.info("Step {} -> best model em: {:.3f}.".format(step, em))
                        copyfile(trans_output, trans_output + '.{}'.format(step))
                        model.save(best_saver, sess, best_model_dir, model_prefix='{}-{}'.format(args.model, step))
                        tf.logging.info("Save best model...")
                    else:
                        tf.logging.info("Step {} -> model acc: {:.3f}, bleu-4: {:.3f}, rouge-L: {:.3f}.".format(
                            step, comp_acc, bleu, rouge))
                        tf.logging.info("Step {} -> model em: {:.3f}.".format(step, em))

            tf.logging.info("Train is end, best Metric: {}".format(best_metric))
示例#2
0
def main(args):
    tf.logging.set_verbosity(tf.logging.INFO)
    model_cls = transformer.Transformer
    args.model = model_cls.get_name()
    params = default_parameters()

    # Import and override parameters
    # Priorities (low -> high):
    # default -> saved -> command
    params = merge_parameters(params, model_cls.get_parameters())
    params = import_params(args.output, args.model, params)
    override_parameters(params, args)

    # Export all parameters and model specific parameters
    export_params(params.output, "params.json", params)
    export_params(params.output, "%s.json" % args.model,
                  collect_params(params, model_cls.get_parameters()))

    #tf.set_random_seed(params.seed)

    # Build Graph
    with tf.Graph().as_default():
        # Build input queue
        features = dataset.get_training_input(params.input, params)

        # features, init_op = cache.cache_features(features, params.update_cycle)
        # Add pre_trained_embedding:
        if params.use_pretrained_embedding:
            _, src_embs = dataset.get_pre_embeddings(params.embeddings[0])
            _, trg_embs = dataset.get_pre_embeddings(params.embeddings[1])
            features['src_embs'] = src_embs
            features['trg_embs'] = trg_embs
            print('Loaded Embeddings!', src_embs.shape, trg_embs.shape)

        # Build model
        initializer = get_initializer(params)
        model = model_cls(params, args.model)

        # Multi-GPU setting
        sharded_losses = parallel.parallel_model(
            model.get_training_func(initializer), features, params.device_list)
        loss = tf.add_n(sharded_losses) / len(sharded_losses)

        # Create global step
        global_step = tf.train.get_or_create_global_step()
        initial_global_step = global_step.assign(0)

        # Print parameters
        all_weights = {v.name: v for v in tf.trainable_variables()}
        total_size = 0

        for v_name in sorted(list(all_weights)):
            v = all_weights[v_name]
            tf.logging.info("%s\tshape    %s", v.name[:-2].ljust(80),
                            str(v.shape).ljust(20))
            v_size = np.prod(np.array(v.shape.as_list())).tolist()
            total_size += v_size
        tf.logging.info("Total trainable variables size: %d", total_size)

        learning_rate = get_learning_rate_decay(params.learning_rate,
                                                global_step, params)
        if params.learning_rate_minimum:
            lr_min = float(params.learning_rate_minimum)
            learning_rate = tf.maximum(learning_rate, tf.to_float(lr_min))

        learning_rate = tf.convert_to_tensor(learning_rate, dtype=tf.float32)
        tf.summary.scalar("learning_rate", learning_rate)

        # Create optimizer
        if params.optimizer == "Adam":
            opt = tf.train.AdamOptimizer(learning_rate,
                                         beta1=params.adam_beta1,
                                         beta2=params.adam_beta2,
                                         epsilon=params.adam_epsilon)
        elif params.optimizer == "LazyAdam":
            opt = tf.contrib.opt.LazyAdamOptimizer(learning_rate,
                                                   beta1=params.adam_beta1,
                                                   beta2=params.adam_beta2,
                                                   epsilon=params.adam_epsilon)
        else:
            raise RuntimeError("Optimizer %s not supported" % params.optimizer)

        loss, ops = optimize.create_train_op(loss, opt, global_step, params)

        restore_op = restore_variables(args.output)

        # Validation
        if params.validation and params.references[0]:
            files = [params.validation] + list(params.references)
            eval_inputs = dataset.sort_and_zip_files(files)
            eval_input_fn = dataset.get_evaluation_input
        else:
            eval_input_fn = None

        # Add hooks
        save_vars = tf.trainable_variables() + [global_step]
        saver = tf.train.Saver(
            var_list=save_vars if params.only_save_trainable else None,
            max_to_keep=params.keep_checkpoint_max,
            sharded=False)
        tf.add_to_collection(tf.GraphKeys.SAVERS, saver)

        train_hooks = [
            tf.train.StopAtStepHook(last_step=params.train_steps),
            #tf.train.StopAtStepHook(num_steps=params.train_steps),
            tf.train.NanTensorHook(loss),
            tf.train.LoggingTensorHook({
                "step": global_step,
                "loss": loss,
            },
                                       every_n_iter=params.print_steps),
            tf.train.CheckpointSaverHook(
                checkpoint_dir=params.output,
                save_secs=params.save_checkpoint_secs or None,
                save_steps=params.save_checkpoint_steps or None,
                saver=saver)
        ]

        config = session_config(params)

        if eval_input_fn is not None:
            train_hooks.append(
                hooks.EvaluationHook(
                    lambda f: beamsearch.create_inference_graph(
                        [model.get_inference_func()], f, params),
                    lambda: eval_input_fn(eval_inputs, params),
                    lambda x: decode_target_ids(x, params),
                    params.output,
                    config,
                    params.keep_top_checkpoint_max,
                    eval_steps_begin=params.eval_steps_begin,
                    eval_secs=params.eval_secs,
                    eval_steps=params.eval_steps))

        def restore_fn(step_context):
            step_context.session.run(restore_op)

        def step_fn(step_context):
            # Bypass hook calls
            return step_context.run_with_hooks(ops)

        # Create session, do not use default CheckpointSaverHook
        with tf.train.MonitoredTrainingSession(checkpoint_dir=params.output,
                                               hooks=train_hooks,
                                               save_checkpoint_secs=None,
                                               config=config) as sess:
            #sess.run(features['source'].eval())
            #sess.run(features['target'].eval())
            # Restore pre-trained variables
            sess.run_step_fn(restore_fn)
            if params.renew_lr == True:
                sess.run(initial_global_step)

            while not sess.should_stop():
                sess.run_step_fn(step_fn)
示例#3
0
def main(args):
    tf.logging.set_verbosity(tf.logging.INFO)
    model_cls = models.get_model(args.model)

    params = default_parameters()

    params = merge_parameters(params, model_cls.get_parameters())

    params = import_params(args.output, args.model, params)
    override_parameters(params, args)

    export_params(params.output, "params.json", params)
    export_params(params.output, "%s.json" % args.model,
                  collect_params(params, model_cls.get_parameters()))

    with tf.Graph().as_default():
        features = dataset.get_training_input(params.input, params)

        update_cycle = params.update_cycle
        features, init_op = cache.cache_features(features, update_cycle)

        initializer = get_initializer(params)
        regularizer = tf.contrib.layers.l1_l2_regularizer(
            scale_l1=params.scale_l1, scale_l2=params.scale_l2)
        model = model_cls(params)
        global_step = tf.train.get_or_create_global_step()

        sharded_losses = parallel.parallel_model(
            model.get_training_func(initializer, regularizer), features,
            params.device_list)

        loss = tf.add_n(sharded_losses) / len(sharded_losses)
        loss = loss + tf.losses.get_regularization_loss()

        all_weights = {v.name: v for v in tf.trainable_variables()}
        total_size = 0

        for v_name in sorted(list(all_weights)):
            v = all_weights[v_name]
            tf.logging.info("%s\tshape    %s", v.name[:-2].ljust(80),
                            str(v.shape).ljust(20))
            v_size = np.prod(np.array(v.shape.as_list())).tolist()
            total_size += v_size
        tf.logging.info("Total trainable variables size: %d", total_size)

        learning_rate = get_learning_rate_decay(params.learning_rate,
                                                global_step, params)
        learning_rate = tf.convert_to_tensor(learning_rate, dtype=tf.float32)
        tf.summary.scalar("learning_rate", learning_rate)

        if params.optimizer == "Adam":
            opt = tf.train.AdamOptimizer(learning_rate,
                                         beta1=params.adam_beta1,
                                         beta2=params.adam_beta2,
                                         epsilon=params.adam_epsilon)
        elif params.optimizer == "LazyAdam":
            opt = tf.contrib.opt.LazyAdamOptimizer(learning_rate,
                                                   beta1=params.adam_beta1,
                                                   beta2=params.adam_beta2,
                                                   epsilon=params.adam_epsilon)
        elif params.optimizer == "SGD":
            opt = tf.train.GradientDescentOptimizer(learning_rate)
        else:
            raise RuntimeError("Optimizer %s not supported" % params.optimizer)

        loss, ops = optimize.create_train_op(loss, opt, global_step, params)
        restore_op = restore_variables(args.checkpoint)

        if params.validation:
            eval_sorted_keys, eval_inputs = dataset.read_eval_input_file(
                params.validation)
            eval_input_fn = dataset.get_predict_input
        else:
            eval_input_fn = None

        save_vars = tf.trainable_variables() + [global_step]
        saver = tf.train.Saver(
            var_list=save_vars if params.only_save_trainable else None,
            max_to_keep=params.keep_checkpoint_max,
            sharded=False)
        tf.add_to_collection(tf.GraphKeys.SAVERS, saver)

        multiplier = tf.convert_to_tensor([update_cycle, 1])

        train_hooks = [
            tf.train.StopAtStepHook(last_step=params.train_steps),
            tf.train.NanTensorHook(loss),
            tf.train.LoggingTensorHook(
                {
                    "step": global_step,
                    "loss": loss,
                    "text": tf.shape(features["text"]) * multiplier,
                    "aspect": tf.shape(features["aspect"]) * multiplier,
                    "polarity": tf.shape(features["polarity"]) * multiplier
                },
                every_n_iter=1),
            tf.train.CheckpointSaverHook(
                checkpoint_dir=params.output,
                save_secs=params.save_checkpoint_secs or None,
                save_steps=params.save_checkpoint_steps or None,
                saver=saver)
        ]

        config = session_config(params)

        if eval_input_fn is not None:
            train_hooks.append(
                hooks.EvaluationHook(
                    lambda f: inference.create_predict_graph([model], f, params
                                                             ),
                    lambda: eval_input_fn(eval_inputs, params),
                    params.output,
                    config,
                    params.keep_top_checkpoint_max,
                    eval_secs=params.eval_secs,
                    eval_steps=params.eval_steps))

        def restore_fn(step_context):
            step_context.session.run(restore_op)

        def step_fn(step_context):
            step_context.session.run([init_op, ops["zero_op"]])
            for i in range(update_cycle - 1):
                step_context.session.run(ops["collect_op"])

            return step_context.run_with_hooks(ops["train_op"])

        with tf.train.MonitoredTrainingSession(checkpoint_dir=params.output,
                                               hooks=train_hooks,
                                               save_checkpoint_secs=None,
                                               config=config) as sess:
            sess.run_step_fn(restore_fn)

            while not sess.should_stop():
                sess.run_step_fn(step_fn)