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))
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)
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)