def main(args): tf.logging.set_verbosity(tf.logging.INFO) model_cls = models.get_model(args.model) 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())) # Build Graph with tf.Graph().as_default(): if not params.record: # Build input queue if params.use_bert and params.bert_emb_path: features = dataset.get_training_input_with_bert( params.input + [params.bert_emb_path], params) else: features = dataset.get_training_input(params.input, params) else: features = record.get_input_features( # ??? os.path.join(params.record, "*train*"), "train", params) # Build model initializer = get_initializer(params) model = model_cls(params) # 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() # 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() # mutiple all dimension size 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) # Create optimizer opt = tf.train.AdamOptimizer(learning_rate, beta1=params.adam_beta1, beta2=params.adam_beta2, epsilon=params.adam_epsilon) if params.update_cycle == 1: train_op = tf.contrib.layers.optimize_loss( name="training", loss=loss, global_step=global_step, learning_rate=learning_rate, clip_gradients=params.clip_grad_norm or None, optimizer=opt, colocate_gradients_with_ops=True) zero_op = tf.no_op("zero_op") collect_op = tf.no_op("collect_op") else: grads_and_vars = opt.compute_gradients( loss, colocate_gradients_with_ops=True) gradients = [item[0] for item in grads_and_vars] variables = [item[1] for item in grads_and_vars] variables = utils.replicate_variables(variables) zero_op = utils.zero_variables(variables) collect_op = utils.collect_gradients(gradients, variables) scale = 1.0 / params.update_cycle gradients, variables = utils.scale_gradients(grads_and_vars, scale) # Gradient clipping avoid greadient explosion!! if isinstance(params.clip_grad_norm or None, float): gradients, _ = tf.clip_by_global_norm(gradients, params.clip_grad_norm) # Update variables grads_and_vars = list(zip(gradients, variables)) with tf.control_dependencies([collect_op]): train_op = opt.apply_gradients(grads_and_vars, global_step) # Validation ''' if params.validation and params.references[0]: files = [params.validation] + list(params.references) eval_inputs = files eval_input_fn = dataset.get_evaluation_input else: print("Don't evaluate") eval_input_fn = None ''' # Add hooks train_hooks = [ tf.train.StopAtStepHook(last_step=params.train_steps), tf.train.NanTensorHook( loss ), # Monitors the loss tensor and stops training if loss is NaN tf.train.LoggingTensorHook( { "step": global_step, "loss": loss, "chars": tf.shape(features["chars"]), "source": tf.shape(features["source"]), #"bert": tf.shape(features["bert"]), "lr": learning_rate }, 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=tf.train.Saver(max_to_keep=params.keep_checkpoint_max, sharded=False)) ] config = session_config(params) ''' if not eval_input_fn is None: train_hooks.append( hooks.EvaluationHook( lambda f: search.create_inference_graph( model.get_evaluation_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_secs=params.eval_secs, eval_steps=params.eval_steps ) ) ''' with tf.train.MonitoredTrainingSession(checkpoint_dir=params.output, hooks=train_hooks, save_checkpoint_secs=None, config=config) as sess: while not sess.should_stop(): utils.session_run(sess, zero_op) for i in range(1, params.update_cycle): utils.session_run(sess, collect_op) sess.run(train_op)
def main(args): if args.distribute: distribute.enable_distributed_training() tf.logging.set_verbosity(tf.logging.INFO) model_cls = models.get_model(args.model) 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 if distribute.rank() == 0: export_params(params.output, "params.json", params) export_params(params.output, "%s.json" % args.model, collect_params(params, model_cls.get_parameters())) # Build Graph with tf.Graph().as_default(): if not params.record: # Build input queue features = dataset.get_training_input(params.input, params) else: features = record.get_input_features( os.path.join(params.record, "*train*"), "train", params) # Build model 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) # Create global step global_step = tf.train.get_or_create_global_step() dtype = tf.float16 if args.half else None # Multi-GPU setting sharded_losses = parallel.parallel_model( model.get_training_func(initializer, regularizer, dtype), features, params.device_list) loss = tf.add_n(sharded_losses) / len(sharded_losses) loss = loss + tf.losses.get_regularization_loss() if distribute.rank() == 0: print_variables() 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("loss", loss) 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) opt = optimizers.MultiStepOptimizer(opt, params.update_cycle) if args.half: opt = optimizers.LossScalingOptimizer(opt, params.loss_scale) # Optimization grads_and_vars = opt.compute_gradients( loss, colocate_gradients_with_ops=True) if params.clip_grad_norm: grads, var_list = list(zip(*grads_and_vars)) grads, _ = tf.clip_by_global_norm(grads, params.clip_grad_norm) grads_and_vars = zip(grads, var_list) train_op = opt.apply_gradients(grads_and_vars, global_step=global_step) # 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 # Hooks train_hooks = [ tf.train.StopAtStepHook(last_step=params.train_steps), tf.train.NanTensorHook(loss), tf.train.LoggingTensorHook( { "step": global_step, "loss": loss, "source": tf.shape(features["source"]), "target": tf.shape(features["target"]) }, every_n_iter=1) ] broadcast_hook = distribute.get_broadcast_hook() if broadcast_hook: train_hooks.append(broadcast_hook) if distribute.rank() == 0: # 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.append( hooks.MultiStepHook(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), step=params.update_cycle)) if eval_input_fn is not None: train_hooks.append( hooks.MultiStepHook(hooks.EvaluationHook( lambda f: inference.create_inference_graph([model], f, params), lambda: eval_input_fn(eval_inputs, params), lambda x: decode_target_ids(x, params), params.output, session_config(params), device_list=params.device_list, max_to_keep=params.keep_top_checkpoint_max, eval_secs=params.eval_secs, eval_steps=params.eval_steps), step=params.update_cycle)) checkpoint_dir = params.output else: checkpoint_dir = None restore_op = restore_variables(args.checkpoint) def restore_fn(step_context): step_context.session.run(restore_op) # Create session, do not use default CheckpointSaverHook with tf.train.MonitoredTrainingSession( checkpoint_dir=checkpoint_dir, hooks=train_hooks, save_checkpoint_secs=None, config=session_config(params)) as sess: # Restore pre-trained variables sess.run_step_fn(restore_fn) while not sess.should_stop(): sess.run(train_op)
def main(args): tf.logging.set_verbosity(tf.logging.INFO) model_cls = models.get_model(args.model) 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())) #import ipdb; ipdb.set_trace() # Build Graph with tf.Graph().as_default(): if not params.record: # Build input queue features = dataset_c2f_4layers.get_training_input_and_c2f_label( params.input, params.c2f_input, params) else: features = record.get_input_features( os.path.join(params.record, "*train*"), "train", params) update_cycle = params.update_cycle features, init_op = cache.cache_features(features, update_cycle) # Build model 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) # Create global step global_step = tf.train.get_or_create_global_step() # Multi-GPU setting sharded_losses = parallel.parallel_model( model.get_training_func(initializer, regularizer), features, params.device_list) if len(sharded_losses) > 1: losses_mle, losses_l1, losses_l2, losses_l3, losses_l4 = sharded_losses loss_mle = tf.add_n(losses_mle) / len(losses_mle) loss_l1 = tf.add_n(losses_l1) / len(losses_l1) loss_l2 = tf.add_n(losses_l2) / len(losses_l2) loss_l3 = tf.add_n(losses_l3) / len(losses_l3) loss_l4 = tf.add_n(losses_l4) / len(losses_l4) else: loss_mle, loss_l1, loss_l2, loss_l3, loss_l4 = sharded_losses[0] loss = loss_mle + (loss_l1 + loss_l2 + loss_l3 + loss_l4 ) / 4.0 + tf.losses.get_regularization_loss() # 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) 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.checkpoint) #import ipdb; ipdb.set_trace() # Validation if params.validation and params.references[0]: files = [params.validation] + list(params.references) eval_inputs = dataset_c2f_4layers.sort_and_zip_files(files) eval_input_fn = dataset_c2f_4layers.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) 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_mle": loss_mle, "loss_l1": loss_l1, "loss_l2": loss_l2, "loss_l3": loss_l3, "loss_l4": loss_l4, "source": tf.shape(features["source"]) * multiplier, "target": tf.shape(features["target"]) * 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_inference_graph([model], 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_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 # 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"]) # 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: # Restore pre-trained variables sess._tf_sess().run(restore_op) while not sess.should_stop(): sess._tf_sess().run([init_op, ops["zero_op"]]) for i in range(update_cycle - 1): sess._tf_sess().run(ops["collect_op"]) sess.run(ops["train_op"])
def main(args): tf.logging.set_verbosity(tf.logging.INFO) model_cls = models.get_model(args.model) 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())) # Build Graph with tf.Graph().as_default(): if not params.record: # Build input queue features = dataset.get_training_input(params.input, params) else: features = record.get_input_features( os.path.join(params.record, "*train*"), "train", params) # Build model initializer = get_initializer(params) model = model_cls(params) if params.MRT: assert params.batch_size == 1 features = mrt_utils.get_MRT(features, params, 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() # 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) learning_rate = tf.convert_to_tensor(learning_rate, dtype=tf.float32) tf.summary.scalar("learning_rate", learning_rate) # Create optimizer opt = tf.train.AdamOptimizer(learning_rate, beta1=params.adam_beta1, beta2=params.adam_beta2, epsilon=params.adam_epsilon) train_op = tf.contrib.layers.optimize_loss( name="training", loss=loss, global_step=global_step, learning_rate=learning_rate, clip_gradients=params.clip_grad_norm or None, optimizer=opt, colocate_gradients_with_ops=True) # 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 train_hooks = [ tf.train.StopAtStepHook(last_step=params.train_steps), tf.train.NanTensorHook(loss), tf.train.LoggingTensorHook( { "step": global_step, "loss": loss, "source": tf.shape(features["source"]), "target": tf.shape(features["target"]) }, 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=tf.train.Saver(max_to_keep=params.keep_checkpoint_max, sharded=False)) ] config = session_config(params) if eval_input_fn is not None: train_hooks.append( hooks.EvaluationHook( lambda f: search.create_inference_graph( model.get_evaluation_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_secs=params.eval_secs, eval_steps=params.eval_steps)) # 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: while not sess.should_stop(): sess.run(train_op)
def main(args): if args.distribute: distribute.enable_distributed_training() tf.logging.set_verbosity(tf.logging.INFO) model_cls = models.get_model(args.model) 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 if not args.distribute or distribute.rank() == 0: export_params(params.output, "params.json", params) export_params(params.output, "%s.json" % args.model, collect_params(params, model_cls.get_parameters())) assert 'r2l' in params.input[2] # Build Graph use_all_devices(params) with tf.Graph().as_default(): if not params.record: # Build input queue features = dataset.abd_get_training_input(params.input, params) else: features = record.get_input_features( os.path.join(params.record, "*train*"), "train", params) update_cycle = params.update_cycle features, init_op = cache.cache_features(features, update_cycle) # Build model 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) # Create global step global_step = tf.train.get_or_create_global_step() dtype = tf.float16 if args.fp16 else None if args.distribute: training_func = model.get_training_func(initializer, regularizer, dtype) loss = training_func(features) else: # Multi-GPU setting sharded_losses = parallel.parallel_model( model.get_training_func(initializer, regularizer, dtype), features, params.device_list) loss = tf.add_n(sharded_losses) / len(sharded_losses) loss = loss + tf.losses.get_regularization_loss() # Print parameters if not args.distribute or distribute.rank() == 0: print_variables() 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) # 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, distribute.all_reduce if args.distribute else None, args.fp16, params) restore_op = restore_variables(args.checkpoint) # 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.abd_get_evaluation_input else: eval_input_fn = None # Add hooks 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, "source": tf.shape(features["source"]) * multiplier, "target": tf.shape(features["target"]) * multiplier }, every_n_iter=1) ] if args.distribute: train_hooks.append(distribute.get_broadcast_hook()) config = session_config(params) if not args.distribute or distribute.rank() == 0: # 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.append( 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)) if eval_input_fn is not None: if not args.distribute or distribute.rank() == 0: train_hooks.append( hooks.EvaluationHook( lambda f: inference.create_inference_graph([model], 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_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 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"]) # Create session, do not use default CheckpointSaverHook if not args.distribute or distribute.rank() == 0: checkpoint_dir = params.output else: checkpoint_dir = None with tf.train.MonitoredTrainingSession(checkpoint_dir=checkpoint_dir, hooks=train_hooks, save_checkpoint_secs=None, config=config) as sess: # Restore pre-trained variables sess.run_step_fn(restore_fn) 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) model_cls = transformer_cache_fixencoder.Transformer 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())) # Build Graph with tf.Graph().as_default(): if not params.record: # Build input queue features = dataset.get_training_input_src_context( params.input, params) else: features = record.get_input_features( os.path.join(params.record, "*train*"), "train", params) features, init_op = cache.cache_features(features, params.update_cycle) # Build model initializer = get_initializer(params) model = model_cls(params) # 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() # 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) 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.checkpoint) restore_trained_encoder_op = restore_encoder_variables( args.thumt_checkpoint) # Validation if params.validation and params.references[0]: files = [params.validation] + list(params.references) eval_inputs = dataset.sort_and_zip_files_catch(files) eval_input_fn = dataset.get_evaluation_input_catch 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.NanTensorHook(loss), tf.train.LoggingTensorHook({ "step": global_step, "loss": loss, }, 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_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_secs=params.eval_secs, eval_steps=params.eval_steps)) def restore_fn(step_context): step_context.session.run(restore_op) step_context.session.run(restore_trained_encoder_op) def step_fn(step_context): # Bypass hook calls step_context.session.run([init_op, ops["zero_op"]]) for i in range(params.update_cycle): step_context.session.run(ops["collect_op"]) step_context.session.run(ops["scale_op"]) # #################################### # # print some unchanged variable # scale = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, # "transformer/encoder/layer_0/self_attention/layer_norm/scale") # # scale = tf.get_variable("transformer/encoder/layer_0/self_attention/layer_norm/scale") # scale = step_context.session.run(scale[0]) # # print(scale) # # #################################### # # print some changed variable # # scale = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, # "transformer/context/head_to_scalar") # # scale = tf.get_variable("transformer/encoder/layer_0/self_attention/layer_norm/scale") # scale = step_context.session.run(scale[0]) # # print(scale) return step_context.run_with_hooks(ops["train_op"]) # 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: # Restore pre-trained variables sess.run_step_fn(restore_fn) 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() # Import and override parameters # Priorities (low -> high): # default -> saved -> command params = merge_parameters(params, model_cls.get_parameters()) params = import_params(args.checkpoint, args.model, params) override_parameters(params, args) # Build Graph with tf.Graph().as_default(): model = model_cls(params) inputs = read_files(args.input) features = get_features(inputs, params) score_fn = model.get_evaluation_func() # scores = score_fn(features, params) # Multi-GPU setting scores = parallel.parallel_model(score_fn, features, params.device_list) sess_creator = tf.train.ChiefSessionCreator( config=session_config(params)) # Load checkpoint tf.logging.info("Loading %s" % args.checkpoint) var_list = tf.train.list_variables(args.checkpoint) values = {} reader = tf.train.load_checkpoint(args.checkpoint) for (name, shape) in var_list: if not name.startswith(model_cls.get_name()): continue tensor = reader.get_tensor(name) values[name] = tensor ops = set_variables(tf.trainable_variables(), values, model_cls.get_name()) assign_op = tf.group(*ops) # Create session fb = 0 with tf.train.MonitoredSession(session_creator=sess_creator) as sess: # Restore variables sess.run(assign_op) fd = tf.gfile.Open(args.rv_file, "w") if params.model_uncertainty: fm = tf.gfile.Open(args.mean_file, 'w') fv = tf.gfile.Open(args.var_file, 'w') fsm = tf.gfile.Open(args.sen_mean, 'w') fsv = tf.gfile.Open(args.sen_var, 'w') fsr = tf.gfile.Open(args.sen_rv, 'w') while not sess.should_stop(): results = sess.run(scores) fb += 1 message = "Finished batch %d" % fb tf.logging.log(tf.logging.INFO, message) if params.model_uncertainty: rv_score = [] mean_score = [] var_score = [] len_score = [] sen_mean = [] sen_var = [] sen_rv = [] for result in results: rv_score.append(result["rv"].tolist()) mean_score.append(result["mean"].tolist()) var_score.append(result["var"].tolist()) len_score.append(result["len"].tolist()) sen_mean.append(result["sen_mean"].tolist()) sen_var.append(result["sen_var"].tolist()) sen_rv.append(result["sen_rv"].tolist()) rv_score = list(itertools.chain(*rv_score)) mean_score = list(itertools.chain(*mean_score)) var_score = list(itertools.chain(*var_score)) len_score = list(itertools.chain(*len_score)) sen_mean = list(itertools.chain(*sen_mean)) sen_var = list(itertools.chain(*sen_var)) sen_rv = list(itertools.chain(*sen_rv)) # mean_score = results["mean"] # var_score = results["var"] # rv_score = results["rv"] # len_score = results["len"].tolist() len_score = map(int, len_score) for i, l in enumerate(len_score): m = mean_score[i][:l] v = var_score[i][:l] r = rv_score[i][:l] for (i_m, i_v, i_r) in zip(m, v, r): # fd.write('%f/%f/%f ' % (i_m, i_v, i_r)) fd.write('%f ' % i_r) fm.write('%f ' % i_m) fv.write('%f ' % i_v) fd.write("\n") fm.write("\n") fv.write("\n") fsm.write("{}\n".format(sen_mean[i])) fsv.write("{}\n".format(sen_var[i])) fsr.write("{}\n".format(sen_rv[i])) else: results = itertools.chain(*results) for value in results: fd.write("%f\n" % value) fd.close()
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)