def test_add_fairness_metrics_baseline_model(self): config = tf_estimator.RunConfig(model_dir=self.model_dir, save_checkpoints_steps=2) feature_columns, _, _, label_column_name = self.load_dataset.get_feature_columns(include_sensitive_columns=True) # pylint: disable=line-too-long estimator = baseline_model.get_estimator( feature_columns=feature_columns, label_column_name=label_column_name, config=config, model_dir=self.model_dir, hidden_units=self.hidden_units, batch_size=self.batch_size) self.assertIsInstance(estimator, tf_estimator.Estimator) # Adds additional fairness metrics to estimator eval_metrics_fn = self.fairness_metrics.create_fairness_metrics_fn() estimator = tf_estimator.add_metrics(estimator, eval_metrics_fn) train_input_fn, test_input_fn = self._get_train_test_input_fn() estimator.train(input_fn=train_input_fn, steps=self.train_steps) eval_results = estimator.evaluate(input_fn=test_input_fn, steps=self.test_steps) self.assertNotEmpty(eval_results) # # Checks if auc metric is computed for all subgroups for subgroup in self.subgroups: self.assertIn('auc subgroup {}'.format(subgroup), eval_results)
def create_estimator(experiment_dir, hparams, decode_length=20): """Creates an estimator with given hyper parameters.""" if FLAGS.worker_gpu > 1: strategy = tf.distribute.MirroredStrategy() else: strategy = None config = tf_estimator.RunConfig(save_checkpoints_steps=1000, save_summary_steps=300, train_distribute=strategy) model_fn = seq2act_estimator.create_model_fn( hparams, seq2act_estimator.compute_additional_loss\ if hparams.use_additional_loss else None, seq2act_estimator.compute_additional_metric\ if hparams.use_additional_loss else None, compute_seq_accuracy=True, decode_length=decode_length) if FLAGS.reference_checkpoint: latest_checkpoint = tf.train.latest_checkpoint( FLAGS.reference_checkpoint) ws = tf_estimator.WarmStartSettings( ckpt_to_initialize_from=latest_checkpoint, vars_to_warm_start=[ "embed_tokens/task_embed_w", "encode_decode/.*", "output_layer/.*" ]) else: ws = None estimator = tf_estimator.Estimator(model_fn=model_fn, model_dir=experiment_dir, config=config, warm_start_from=ws) return estimator
def test_create_train_and_eval_specs(self): """Tests that `TrainSpec` and `EvalSpec` is created correctly.""" run_config = tf_estimator.RunConfig() hparams = model_hparams.create_hparams( hparams_overrides='load_pretrained=false') pipeline_config_path = get_pipeline_config_path(MODEL_NAME_FOR_TEST) train_steps = 20 train_and_eval_dict = model_lib.create_estimator_and_inputs( run_config, hparams, pipeline_config_path, train_steps=train_steps) train_input_fn = train_and_eval_dict['train_input_fn'] eval_input_fns = train_and_eval_dict['eval_input_fns'] eval_on_train_input_fn = train_and_eval_dict['eval_on_train_input_fn'] predict_input_fn = train_and_eval_dict['predict_input_fn'] train_steps = train_and_eval_dict['train_steps'] train_spec, eval_specs = model_lib.create_train_and_eval_specs( train_input_fn, eval_input_fns, eval_on_train_input_fn, predict_input_fn, train_steps, eval_on_train_data=True, final_exporter_name='exporter', eval_spec_names=['holdout']) self.assertEqual(train_steps, train_spec.max_steps) self.assertEqual(2, len(eval_specs)) self.assertEqual(None, eval_specs[0].steps) self.assertEqual('holdout', eval_specs[0].name) self.assertEqual('exporter', eval_specs[0].exporters[0].name) self.assertEqual(None, eval_specs[1].steps) self.assertEqual('eval_on_train', eval_specs[1].name)
def create_estimator(working_dir, model_fn, keep_checkpoint_max=20, iterations_per_loop=320, warmstart=None): """Create a TF estimator. Used when not using TPU. Args: working_dir: working directory for loading the model. model_fn: an estimator model function. keep_checkpoint_max: the maximum number of checkpoints to save in checkpoint directory. iterations_per_loop: number of steps to run on TPU before outfeeding metrics to the CPU. If the number of iterations in the loop would exceed the number of train steps, the loop will exit before reaching --iterations_per_loop. The larger this value is, the higher the utilization on the TPU. For CPU-only training, this flag is equal to `num_epochs * num_minibatches`. warmstart: if not None, warm start the estimator from an existing checkpoint. """ run_config = tf_estimator.RunConfig( save_checkpoints_steps=iterations_per_loop, save_summary_steps=iterations_per_loop, keep_checkpoint_max=keep_checkpoint_max) if warmstart is not None: return tf_estimator.Estimator(model_fn, model_dir=working_dir, config=run_config, warm_start_from=warmstart) else: return tf_estimator.Estimator(model_fn, model_dir=working_dir, config=run_config)
def train_and_eval(): """Train and evaluate a model.""" save_summary_steps = FLAGS.save_summaries_steps save_checkpoints_steps = FLAGS.save_checkpoints_steps log_step_count = FLAGS.log_step_count config = tf_estimator.RunConfig( save_summary_steps=save_summary_steps, save_checkpoints_steps=save_checkpoints_steps, log_step_count_steps=log_step_count, keep_checkpoint_max=None) params = {'dummy': 0} estimator = tf_estimator.Estimator(model_fn=model_fn, model_dir=FLAGS.checkpoint_dir, config=config, params=params) train_spec = tf_estimator.TrainSpec(input_fn=train_input_fn, max_steps=FLAGS.train_steps) eval_spec = tf_estimator.EvalSpec(input_fn=eval_input_fn, start_delay_secs=60, steps=FLAGS.eval_examples, throttle_secs=60) tf_estimator.train_and_evaluate(estimator, train_spec, eval_spec)
def run(): """Runs train_and_evaluate.""" hparams_filename = os.path.join(FLAGS.model_dir, 'hparams.json') if FLAGS.is_chief: gfile.MakeDirs(FLAGS.model_dir) hparams = core.read_hparams(FLAGS.hparams, get_hparams()) core.write_hparams(hparams, hparams_filename) # Always load HParams from model_dir. hparams = core.wait_for_hparams(hparams_filename, get_hparams()) grammar = grammar_utils.load_grammar(grammar_path=hparams.grammar_path) estimator = tf_estimator.Estimator( model_fn=functools.partial(model_fn, grammar=grammar), params=hparams, config=tf_estimator.RunConfig( save_checkpoints_secs=hparams.save_checkpoints_secs, keep_checkpoint_max=hparams.keep_checkpoint_max)) train_spec = tf_estimator.TrainSpec(input_fn=functools.partial( input_ops.input_fn, input_pattern=hparams.train_pattern, grammar=grammar), max_steps=hparams.train_steps) # NOTE(leeley): The SavedModel will be stored under the # tf.saved_model.tag_constants.SERVING tag. latest_exporter = tf_estimator.LatestExporter( name='latest_exported_model', serving_input_receiver_fn=functools.partial( input_ops.serving_input_receiver_fn, params=hparams, num_production_rules=grammar.num_production_rules), exports_to_keep=hparams.exports_to_keep) eval_hooks = [] if hparams.num_expressions_per_condition > 0: eval_hooks.append( metrics.GenerationWithLeadingPowersHook( generation_leading_powers_abs_sums=core.hparams_list_value( hparams.generation_leading_powers_abs_sums), num_expressions_per_condition=hparams. num_expressions_per_condition, max_length=hparams.max_length, grammar=grammar)) eval_spec = tf_estimator.EvalSpec( input_fn=functools.partial(input_ops.input_fn, input_pattern=hparams.tune_pattern, grammar=grammar), steps=hparams.eval_steps, exporters=latest_exporter, start_delay_secs=hparams.start_delay_secs, throttle_secs=hparams.throttle_secs, hooks=eval_hooks) tf_estimator.train_and_evaluate(estimator, train_spec, eval_spec)
def run_local_training(losses_fn, input_fn, trainer_params_overrides, model_params, vars_to_restore_fn=None): """Run a simple single-mechine traing loop. Args: losses_fn: A callable that receives two arguments, `features` and `params`, both are dictionaries, and returns a dictionary whose values are the losses. Their sum is the total loss to be minimized. input_fn: A callable that complies with tf.Estimtor's definition of input_fn. trainer_params_overrides: A dictionary or a ParameterContainer with overrides for the default values in TRAINER_PARAMS above. model_params: A ParameterContainer that will be passed to the model (i. e. to losses_fn and input_fn). vars_to_restore_fn: A callable that receives no arguments. When called, expected to provide a dictionary that maps the checkpoint name of each variable to the respective variable object. This dictionary will be used as `var_list` in a Saver object used for initializing from the checkpoint at trainer_params.init_ckpt. If None, the default saver will be used. """ trainer_params = ParameterContainer.from_defaults_and_overrides( TRAINER_PARAMS, trainer_params_overrides, is_strict=True) run_config_params = { 'model_dir': trainer_params.model_dir, 'save_summary_steps': 50, 'keep_checkpoint_every_n_hours': trainer_params.keep_checkpoint_every_n_hours, 'log_step_count_steps': 50, } logging.info( 'Estimators run config parameters:\n%s', json.dumps(run_config_params, indent=2, sort_keys=True, default=str)) run_config = tf_estimator.RunConfig(**run_config_params) def estimator_spec_fn(features, labels, mode, params): del labels # unused return _build_estimator_spec(losses_fn(features, mode, params), trainer_params=trainer_params, mode=mode, use_tpu=False) init_hook = InitFromCheckpointHook(trainer_params.model_dir, trainer_params.init_ckpt, vars_to_restore_fn) estimator = tf_estimator.Estimator(model_fn=estimator_spec_fn, config=run_config, params=model_params.as_dict()) estimator.train(input_fn=input_fn, max_steps=trainer_params.max_steps, hooks=[init_hook])
def testEndToEnd(self): params = imagenet_params params['output_dir'] = '/tmp/' params['batch_size'] = 2 params['num_train_steps'] = 1 params['eval_steps'] = 1 params['threshold'] = 80. params['data_format'] = 'channels_last' mean_stats = [0.485, 0.456, 0.406] std_stats = [0.229, 0.224, 0.225] update_params = { 'mean_rgb': mean_stats, 'stddev_rgb': std_stats, 'lr_schedule': [ # (multiplier, epoch to start) tuples (1.0, 5), (0.1, 30), (0.01, 60), (0.001, 80) ], 'momentum': 0.9, 'data_format': 'channels_last' } params.update(update_params) dataset_ = data_input.DataIterator(mode=FLAGS.mode, data_directory='', saliency_method='ig_smooth_2', transformation='modified_image', threshold=params['threshold'], keep_information=False, use_squared_value=True, mean_stats=mean_stats, std_stats=std_stats, test_small_sample=True, num_cores=FLAGS.num_cores) images, labels = dataset_.input_fn(params) self.assertEqual(images.shape.as_list(), [2, 224, 224, 3]) self.assertEqual(labels.shape.as_list(), [ 2, ]) run_config = tf_estimator.RunConfig( model_dir=FLAGS.dest_dir, save_checkpoints_steps=FLAGS.steps_per_checkpoint) classifier = tf_estimator.Estimator(model_fn=resnet_model_fn, model_dir=FLAGS.dest_dir, params=params, config=run_config) classifier.train(input_fn=dataset_.input_fn, max_steps=1) tf.logging.info('finished training.')
def test_experiment(self): """Tests that the `Experiment` object is constructed correctly.""" run_config = tf_estimator.RunConfig() hparams = model_hparams.create_hparams( hparams_overrides='load_pretrained=false') pipeline_config_path = get_pipeline_config_path(MODEL_NAME_FOR_TEST) experiment = model_lib.populate_experiment(run_config, hparams, pipeline_config_path, train_steps=10, eval_steps=20) self.assertEqual(10, experiment.train_steps) self.assertEqual(None, experiment.eval_steps)
def initiate_task_helper(model_params, ckpt_directory=None, pruning_params=None): """Get all predictions for eval. Args: model_params: ckpt_directory: model checkpoint directory containing event file pruning_params: Returns: pd.DataFrame containing metrics from event file """ if model_params["task"] != "imagenet_training": classifier = tf_estimator.Estimator( model_fn=model_fn_w_pruning, params=model_params) if model_params["task"] in ["imagenet_predictions"]: predictions = classifier.predict( input_fn=data_input.input_fn, checkpoint_path=ckpt_directory) return predictions if model_params["task"] in [ "robustness_imagenet_a", "robustness_imagenet_c", "robustness_pie", "imagenet_eval", "ckpt_prediction" ]: eval_steps = model_params["num_eval_images"] // model_params["batch_size"] tf.logging.info("start computing eval metrics...") classifier = tf_estimator.Estimator( model_fn=model_fn_w_pruning, params=model_params) evaluation_metrics = classifier.evaluate( input_fn=data_input.input_fn, steps=eval_steps, checkpoint_path=ckpt_directory) tf.logging.info("finished per class accuracy eval.") return evaluation_metrics else: model_params["pruning_dict"] = pruning_params run_config = tf_estimator.RunConfig( save_summary_steps=300, save_checkpoints_steps=1000, log_step_count_steps=100) classifier = tf_estimator.Estimator( model_fn=model_fn_w_pruning, config=run_config, params=model_params) tf.logging.info("start training...") classifier.train( input_fn=data_input.input_fn, max_steps=model_params["num_train_steps"]) tf.logging.info("finished training.")
def main(argv): del argv # Unused. if FLAGS.output_dir is None: raise ValueError("`output_dir` must be defined") if FLAGS.delete_existing and tf.gfile.Exists(FLAGS.output_dir): tf.logging.warn("Deleting old log directory at {}".format( FLAGS.output_dir)) tf.gfile.DeleteRecursively(FLAGS.output_dir) tf.gfile.MakeDirs(FLAGS.output_dir) print("Logging to {}".format(FLAGS.output_dir)) # Load the training or test split of the Celeb-A filenames. if FLAGS.celeba_dir is None: raise ValueError("`celeba_dir` must be defined") celeba_dataset_path = \ os.path.join(FLAGS.celeba_dir, "Img/img_align_celeba/") celeba_partition_path = \ os.path.join(FLAGS.celeba_dir, "Eval/list_eval_partition.txt") with open(celeba_partition_path, "r") as fid: partition = fid.readlines() filenames, splits = zip(*[x.split() for x in partition]) filenames = np.array( [os.path.join(celeba_dataset_path, f) for f in filenames]) splits = np.array([int(x) for x in splits]) with tf.Graph().as_default(): train_input_fn = prep_dataset_fn(filenames, splits, is_training=True) eval_input_fn = prep_dataset_fn(filenames, splits, is_training=False) estimator = tf_estimator.Estimator( model_fn, config=tf_estimator.RunConfig( model_dir=FLAGS.output_dir, save_checkpoints_steps=FLAGS.viz_steps, ), ) train_spec = tf_estimator.TrainSpec(input_fn=train_input_fn, max_steps=FLAGS.max_steps) # Sad ugly hack here. Setting steps=None should go through all of the # validation set, but doesn't seem to, so I'm doing it manually. eval_spec = tf_estimator.EvalSpec(input_fn=eval_input_fn, steps=len(filenames[splits == 1]) // FLAGS.batch_size, start_delay_secs=0, throttle_secs=0) for _ in range(FLAGS.max_steps // FLAGS.viz_steps): tf_estimator.train_and_evaluate(estimator, train_spec, eval_spec)
def run_std_server(self): """Starts a TensorFlow server and joins the serving thread. Typically used for parameter servers. Raises: ValueError: if not enough information is available in the estimator's config to create a server. """ config = tf_estimator.RunConfig() server = tf.train.Server(config.cluster_spec, job_name=config.task_type, task_index=config.task_id, protocol=config.protocol) server.join()
def test_create_baseline_estimator_without_demographics(self): config = tf_estimator.RunConfig(model_dir=self.model_dir, save_checkpoints_steps=2) feature_columns, _, _, label_column_name = self.load_dataset.get_feature_columns(include_sensitive_columns=False) # pylint: disable=line-too-long estimator = baseline_model.get_estimator( feature_columns=feature_columns, label_column_name=label_column_name, config=config, model_dir=self.model_dir, hidden_units=self.hidden_units, batch_size=self.batch_size, learning_rate=0.01, optimizer='Adagrad', activation=tf.nn.relu) self.assertIsInstance(estimator, tf_estimator.Estimator)
def main(_): cpu = os.cpu_count() tf_config = _tf_config(flags) #1 # 分布式需要 TF_CONFIG 环境变量 os.environ['TF_CONFIG'] = json.dumps(tf_config) #2 session_config = ConfigProto(device_count={'CPU': cpu}, inter_op_parallelism_threads=cpu // 2, intra_op_parallelism_threads=cpu // 2, device_filters=flags.device_filters, allow_soft_placement=True) strategy = experimental.ParameterServerStrategy() run_config = estimator.RunConfig( **{ 'save_summary_steps': 100, 'save_checkpoints_steps': 1000, 'keep_checkpoint_max': 10, 'log_step_count_steps': 100, 'train_distribute': strategy, 'eval_distribute': strategy, }).replace(session_config=session_config) model = estimator.Estimator( model_fn=model_fn, model_dir='/home/axing/din/checkpoints/din', #实际应用中是分布式文件系统 config=run_config, params={ 'tf_config': tf_config, 'decay_rate': 0.9, 'decay_steps': 10000, 'learning_rate': 0.1 }) train_spec = estimator.TrainSpec( input_fn=lambda: input_fn(mode='train', num_workers=flags.num_workers, worker_index=flags.worker_index, pattern='/home/axing/din/dataset/*'), #3 max_steps=1000 #4 ) # 这里就假设验证集和训练集地址一样了,实际应用中是肯定不一样的。 eval_spec = estimator.EvalSpec( input_fn=lambda: input_fn(mode='eval', pattern='/home/axing/din/dataset/*'), steps=100, # 每次验证 100 个 batch size 的数据 throttle_secs=60 # 每隔至少 60 秒验证一次 ) estimator.train_and_evaluate(model, train_spec, eval_spec)
def test_create_estimator_with_default_train_eval_steps(self): """Tests that number of train/eval defaults to config values.""" run_config = tf_estimator.RunConfig() hparams = model_hparams.create_hparams( hparams_overrides='load_pretrained=false') pipeline_config_path = get_pipeline_config_path(MODEL_NAME_FOR_TEST) configs = config_util.get_configs_from_pipeline_file( pipeline_config_path) config_train_steps = configs['train_config'].num_steps train_and_eval_dict = model_lib.create_estimator_and_inputs( run_config, hparams, pipeline_config_path) estimator = train_and_eval_dict['estimator'] train_steps = train_and_eval_dict['train_steps'] self.assertIsInstance(estimator, tf_estimator.Estimator) self.assertEqual(config_train_steps, train_steps)
def test_create_estimator_and_inputs(self): """Tests that Estimator and input function are constructed correctly.""" run_config = tf_estimator.RunConfig() hparams = model_hparams.create_hparams( hparams_overrides='load_pretrained=false') pipeline_config_path = get_pipeline_config_path(MODEL_NAME_FOR_TEST) train_steps = 20 train_and_eval_dict = model_lib.create_estimator_and_inputs( run_config, hparams, pipeline_config_path, train_steps=train_steps) estimator = train_and_eval_dict['estimator'] train_steps = train_and_eval_dict['train_steps'] self.assertIsInstance(estimator, tf_estimator.Estimator) self.assertEqual(20, train_steps) self.assertIn('train_input_fn', train_and_eval_dict) self.assertIn('eval_input_fns', train_and_eval_dict) self.assertIn('eval_on_train_input_fn', train_and_eval_dict)
def main(unused_argv): flags.mark_flag_as_required('model_dir') flags.mark_flag_as_required('pipeline_config_path') config = tf_estimator.RunConfig(model_dir=FLAGS.model_dir) train_and_eval_dict = model_lib.create_estimator_and_inputs( run_config=config, pipeline_config_path=FLAGS.pipeline_config_path, train_steps=FLAGS.num_train_steps, sample_1_of_n_eval_examples=FLAGS.sample_1_of_n_eval_examples, sample_1_of_n_eval_on_train_examples=( FLAGS.sample_1_of_n_eval_on_train_examples)) estimator = train_and_eval_dict['estimator'] train_input_fn = train_and_eval_dict['train_input_fn'] eval_input_fns = train_and_eval_dict['eval_input_fns'] eval_on_train_input_fn = train_and_eval_dict['eval_on_train_input_fn'] predict_input_fn = train_and_eval_dict['predict_input_fn'] train_steps = train_and_eval_dict['train_steps'] if FLAGS.checkpoint_dir: if FLAGS.eval_training_data: name = 'training_data' input_fn = eval_on_train_input_fn else: name = 'validation_data' # The first eval input will be evaluated. input_fn = eval_input_fns[0] if FLAGS.run_once: estimator.evaluate(input_fn, steps=None, checkpoint_path=tf.train.latest_checkpoint( FLAGS.checkpoint_dir)) else: model_lib.continuous_eval(estimator, FLAGS.checkpoint_dir, input_fn, train_steps, name, FLAGS.max_eval_retries) else: train_spec, eval_specs = model_lib.create_train_and_eval_specs( train_input_fn, eval_input_fns, eval_on_train_input_fn, predict_input_fn, train_steps, eval_on_train_data=False) # Currently only a single Eval Spec is allowed. tf_estimator.train_and_evaluate(estimator, train_spec, eval_specs[0])
def get_estimator(self): """Obtain estimator for the working directory. Returns: an (TPU/non-TPU) estimator. """ if self._use_tpu: return tf_utils.get_tpu_estimator(self._checkpoint_dir, self._model_fn) run_config = tf_estimator.RunConfig( save_summary_steps=self._iterations_per_loop, save_checkpoints_steps=self._iterations_per_loop, keep_checkpoint_max=self._keep_checkpoint_max) return tf_estimator.Estimator(self._model_fn, model_dir=self._checkpoint_dir, config=run_config)
def test_global_steps_baseline_model(self): config = tf_estimator.RunConfig(model_dir=self.model_dir, save_checkpoints_steps=2) feature_columns, _, _, label_column_name = self.load_dataset.get_feature_columns(include_sensitive_columns=True) # pylint: disable=line-too-long estimator = baseline_model.get_estimator( feature_columns=feature_columns, label_column_name=label_column_name, config=config, model_dir=self.model_dir, hidden_units=self.hidden_units, batch_size=self.batch_size) self.assertIsInstance(estimator, tf_estimator.Estimator) train_input_fn, test_input_fn = self._get_train_test_input_fn() estimator.train(input_fn=train_input_fn, steps=self.train_steps) eval_results = estimator.evaluate(input_fn=test_input_fn, steps=self.test_steps) # Checks if global step has reached specified number of train_steps self.assertIn('global_step', eval_results) self.assertEqual(eval_results['global_step'], self.train_steps)
def test_create_adversarial_reweighting_estimator_without_demographics(self): config = tf_estimator.RunConfig(model_dir=self.model_dir, save_checkpoints_steps=2) feature_columns, _, _, label_column_name = self.load_dataset.get_feature_columns(include_sensitive_columns=False) # pylint: disable=line-too-long estimator = adversarial_reweighting_model.get_estimator( feature_columns=feature_columns, label_column_name=label_column_name, config=config, model_dir=self.model_dir, primary_hidden_units=self.primary_hidden_units, batch_size=self.batch_size, pretrain_steps=self.pretrain_steps, primary_learning_rate=0.01, adversary_learning_rate=0.01, optimizer='Adagrad', activation=tf.nn.relu, adversary_loss_type='ce_loss', adversary_include_label=True, upweight_positive_instance_only=False) self.assertIsInstance(estimator, tf_estimator.Estimator)
def _get_estimator(self): config = tf_estimator.RunConfig(model_dir=self.model_dir, save_checkpoints_steps=1) feature_columns, _, _, label_column_name = ( self.load_dataset.get_feature_columns( include_sensitive_columns=True)) estimator = adversarial_reweighting_model.get_estimator( feature_columns=feature_columns, label_column_name=label_column_name, config=config, model_dir=self.model_dir, primary_hidden_units=self.primary_hidden_units, batch_size=self.batch_size, pretrain_steps=self.pretrain_steps, primary_learning_rate=0.01, adversary_learning_rate=0.01, optimizer='Adagrad', activation=tf.nn.relu, adversary_loss_type='ce_loss', adversary_include_label=True, upweight_positive_instance_only=False) return estimator
def main(_): inference_fn = network.inference hparams = contrib_training.HParams(learning_rate=FLAGS.learning_rate) model_fn = estimator.create_model_fn(inference_fn, hparams) config = tf_estimator.RunConfig(FLAGS.model_dir) tf_estimator = tf_estimator.Estimator(model_fn=model_fn, config=config) train_dataset_fn = dataset.create_dataset_fn(FLAGS.train_pattern, height=FLAGS.image_size, width=FLAGS.image_size, batch_size=FLAGS.batch_size) eval_dataset_fn = dataset.create_dataset_fn(FLAGS.test_pattern, height=FLAGS.image_size, width=FLAGS.image_size, batch_size=FLAGS.batch_size) train_spec, eval_spec = estimator.create_train_and_eval_specs( train_dataset_fn, eval_dataset_fn) tf.logging.set_verbosity(tf.logging.INFO) tf_estimator.train_and_evaluate(tf_estimator, train_spec, eval_spec)
def _make_estimator(hparams, label_vocab, output_dir): """Create a tf.estimator.Estimator. Args: hparams: tf.contrib.training.HParams. label_vocab: list of string. output_dir: str. Path to save checkpoints. Returns: tf.estimator.Estimator. """ model_fn = protein_model.make_model_fn(label_vocab=label_vocab, hparams=hparams) run_config = tf_estimator.RunConfig(model_dir=output_dir) estimator = tf_estimator.Estimator( model_fn=model_fn, params=hparams, config=run_config, ) return estimator
def test_eval_results_baseline_model(self): config = tf_estimator.RunConfig(model_dir=self.model_dir, save_checkpoints_steps=2) feature_columns, _, _, label_column_name = self.load_dataset.get_feature_columns(include_sensitive_columns=True) # pylint: disable=line-too-long estimator = baseline_model.get_estimator( feature_columns=feature_columns, label_column_name=label_column_name, config=config, model_dir=self.model_dir, hidden_units=self.hidden_units, batch_size=self.batch_size) self.assertIsInstance(estimator, tf_estimator.Estimator) train_input_fn, test_input_fn = self._get_train_test_input_fn() estimator.train(input_fn=train_input_fn, steps=self.train_steps) eval_results = estimator.evaluate(input_fn=test_input_fn, steps=self.test_steps) self.assertNotEmpty(eval_results) # # Checks if all tp,tn,fp,fn keys are present in eval_results dictionary self.assertIn('auc', eval_results) self.assertIn('fp', eval_results) self.assertIn('fn', eval_results) self.assertIn('tp', eval_results) self.assertIn('tn', eval_results)
def main(_): tf.logging.set_verbosity(tf.logging.INFO) train_mode = FLAGS.train_mode ############################################################################ # Load the dual_encoder_config_file file. ############################################################################ if tf.gfile.Exists(FLAGS.dual_encoder_config_file): exp_config = utils.load_config_from_file( FLAGS.dual_encoder_config_file, experiment_config_pb2.DualEncoderConfig()) else: raise ValueError("dual_encoder_config: {} not found!".format( FLAGS.dual_encoder_config_file)) tf.logging.info(">>>> final dual_encoder_config:\n {}".format(exp_config)) tf.gfile.MakeDirs(FLAGS.output_dir) ############################################################################ # Save/copy the configuration file. ############################################################################ configs_dir = os.path.join(FLAGS.output_dir, "configs") tf.gfile.MakeDirs(configs_dir) tf.gfile.MakeDirs(FLAGS.output_dir) with tf.gfile.Open(os.path.join(configs_dir, "dual_encoder_config.pbtxt"), "w") as fout: print(exp_config, file=fout) # Write bert_config.json and doc_bert_config.json. tf.gfile.Copy(exp_config.encoder_config.bert_config_file, os.path.join(configs_dir, "bert_config.json"), overwrite=True) tf.gfile.Copy(exp_config.encoder_config.doc_bert_config_file, os.path.join(configs_dir, "doc_bert_config.json"), overwrite=True) # Write vocab file(s). tf.gfile.Copy(exp_config.encoder_config.vocab_file, os.path.join(configs_dir, "vocab.txt"), overwrite=True) # Save other important parameters as a json file. hparams = { "dual_encoder_config_file": FLAGS.dual_encoder_config_file, "output_dir": FLAGS.output_dir, "schedule": FLAGS.schedule, "debugging": FLAGS.debugging, "learning_rate": FLAGS.learning_rate, "num_warmup_steps": FLAGS.num_warmup_steps, "num_train_steps": FLAGS.num_train_steps, "num_tpu_cores": FLAGS.num_tpu_cores } with tf.gfile.Open(os.path.join(configs_dir, "hparams.json"), "w") as fout: json.dump(hparams, fout) tf.logging.info(">>>> saved hparams.json:\n {}".format(hparams)) ############################################################################ # Run the train/eval/predict/export process based on the schedule. ############################################################################ max_seq_length_actual, max_predictions_per_seq_actual = \ utils.get_actual_max_seq_len(exp_config.encoder_config.model_name, exp_config.encoder_config.max_doc_length_by_sentence, exp_config.encoder_config.max_sent_length_by_word, exp_config.encoder_config.max_predictions_per_seq) # Prepare input for train and eval. input_files = [] for input_pattern in exp_config.train_eval_config.input_file_for_train.split( ","): input_files.extend(tf.gfile.Glob(input_pattern)) input_file_num = 0 tf.logging.info("*** Input Files ***") for input_file in input_files: tf.logging.info(" %s" % input_file) input_file_num += 1 if input_file_num > 10: break tf.logging.info("train input_files[0:10]: %s " % "\n".join(input_files[0:10])) eval_files = [] if exp_config.train_eval_config.eval_with_eval_data: eval_files = [] for input_pattern in exp_config.train_eval_config.input_file_for_eval.split( ","): eval_files.extend(tf.gfile.Glob(input_pattern)) else: eval_files = input_files input_fn_builder = input_fns.input_fn_builder # Prepare the input functions. # Drop_remainder = True during training to maintain fixed batch size. train_input_fn = input_fn_builder( input_files=input_files, is_training=True, drop_remainder=True, max_seq_length=max_seq_length_actual, max_predictions_per_seq=max_predictions_per_seq_actual, num_cpu_threads=4, batch_size=exp_config.train_eval_config.train_batch_size, ) eval_drop_remainder = True if FLAGS.use_tpu else False eval_input_fn = input_fn_builder( input_files=eval_files, max_seq_length=max_seq_length_actual, max_predictions_per_seq=max_predictions_per_seq_actual, is_training=False, drop_remainder=eval_drop_remainder, batch_size=exp_config.train_eval_config.eval_batch_size) predict_input_fn = input_fn_builder( input_files=eval_files, max_seq_length=max_seq_length_actual, max_predictions_per_seq=max_predictions_per_seq_actual, is_training=False, drop_remainder=eval_drop_remainder, batch_size=exp_config.train_eval_config.predict_batch_size, is_prediction=True) # Build and run the model. tpu_cluster_resolver = None if FLAGS.use_tpu and FLAGS.tpu_name: tpu_cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver( FLAGS.tpu_name, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project) is_per_host = tf_estimator.tpu.InputPipelineConfig.PER_HOST_V2 run_config = tf_estimator.tpu.RunConfig( cluster=tpu_cluster_resolver, master=FLAGS.master, model_dir=FLAGS.output_dir, save_checkpoints_steps=exp_config.train_eval_config. save_checkpoints_steps, tpu_config=tf_estimator.tpu.TPUConfig( iterations_per_loop=exp_config.train_eval_config. iterations_per_loop, num_shards=FLAGS.num_tpu_cores, per_host_input_for_training=is_per_host)) model_fn = smith_modeling.model_fn_builder( dual_encoder_config=exp_config, train_mode=FLAGS.train_mode, learning_rate=FLAGS.learning_rate, num_train_steps=FLAGS.num_train_steps, num_warmup_steps=FLAGS.num_warmup_steps, use_tpu=FLAGS.use_tpu, use_one_hot_embeddings=FLAGS.use_tpu, debugging=FLAGS.debugging) # If TPU is not available, this will fall back to normal Estimator on CPU # or GPU. The batch size for eval and predict is the same. estimator = tf_estimator.tpu.TPUEstimator( use_tpu=FLAGS.use_tpu, model_fn=model_fn, config=run_config, train_batch_size=exp_config.train_eval_config.train_batch_size, eval_batch_size=exp_config.train_eval_config.eval_batch_size, predict_batch_size=exp_config.train_eval_config.predict_batch_size) if FLAGS.schedule == "train": tf.logging.info("***** Running training *****") tf.logging.info(" Batch size = %d", exp_config.train_eval_config.train_batch_size) estimator.train(input_fn=train_input_fn, max_steps=FLAGS.num_train_steps) elif FLAGS.schedule == "continuous_eval": tf.logging.info("***** Running continuous evaluation *****") tf.logging.info(" Batch size = %d", exp_config.train_eval_config.eval_batch_size) # checkpoints_iterator blocks until a new checkpoint appears. for ckpt in tf.train.checkpoints_iterator(estimator.model_dir): try: # Estimator automatically loads and evaluates the latest checkpoint. result = estimator.evaluate( input_fn=eval_input_fn, steps=exp_config.train_eval_config.max_eval_steps) tf.logging.info("***** Eval results for %s *****", ckpt) for key, value in result.items(): tf.logging.info(" %s = %s", key, str(value)) except tf.errors.NotFoundError: # Checkpoint might get garbage collected before the eval can run. tf.logging.error("Checkpoint path '%s' no longer exists.", ckpt) elif FLAGS.schedule == "predict": # Load the model checkpoint and run the prediction process # to get the predicted scores and labels. The batch size is the same with # the eval batch size. For more options, refer to # https://www.tensorflow.org/api_docs/python/tf/compat/v1/estimator/tpu/TPUEstimator#predict tf.logging.info("***** Running prediction with ckpt {} *****".format( exp_config.encoder_config.predict_checkpoint)) tf.logging.info(" Batch size = %d", exp_config.train_eval_config.eval_batch_size) output_predict_file = os.path.join(FLAGS.output_dir, "prediction_results.json") # Output the prediction results in json format. pred_res_list = [] with tf.gfile.GFile(output_predict_file, "w") as writer: written_line_index = 0 tf.logging.info("***** Predict results *****") for result in estimator.predict(input_fn=predict_input_fn, checkpoint_path=exp_config. encoder_config.predict_checkpoint, yield_single_examples=True): if (exp_config.encoder_config.model_name == constants.MODEL_NAME_SMITH_DUAL_ENCODER): pred_item_dict = utils.get_pred_res_list_item_smith_de( result) else: raise ValueError("Unsupported model name: %s" % exp_config.encoder_config.model_name) pred_res_list.append(pred_item_dict) written_line_index += 1 if written_line_index % 500 == 0: tf.logging.info( "Current written_line_index: {} *****".format( written_line_index)) tf.logging.info("***** Finished prediction for %d examples *****", written_line_index) tf.logging.info("***** Output prediction results into %s *****", output_predict_file) json.dump(pred_res_list, writer) elif FLAGS.schedule == "export": run_config = tf_estimator.RunConfig( model_dir=FLAGS.output_dir, save_checkpoints_steps=exp_config.train_eval_config. save_checkpoints_steps) estimator = tf_estimator.Estimator(model_fn=model_fn, config=run_config) export_dir_base = os.path.join(FLAGS.output_dir, "export/") tf.logging.info( "***** Export the prediction checkpoint to the folder {} *****". format(export_dir_base)) tf.gfile.MakeDirs(export_dir_base) estimator.export_saved_model( export_dir_base=export_dir_base, assets_extra={"vocab.txt": exp_config.encoder_config.vocab_file}, serving_input_receiver_fn=input_fns.make_serving_input_example_fn( max_seq_length=max_seq_length_actual, max_predictions_per_seq=max_predictions_per_seq_actual), checkpoint_path=exp_config.encoder_config.predict_checkpoint) else: raise ValueError("Unsupported schedule : %s" % FLAGS.schedule)
def main(_): tf.logging.set_verbosity(tf.logging.INFO) if FLAGS.model == "seq2seq": assert FLAGS.rnn_cell == "lstm" assert FLAGS.att_type != "hyper" if FLAGS.model == "hypernet" and FLAGS.rank != FLAGS.decoder_dim: print("WARNING: recommended rank value: decoder_dim.") if FLAGS.att_neighbor: assert FLAGS.neighbor_dim == FLAGS.encoder_dim or FLAGS.att_type == "my" if FLAGS.use_copy or FLAGS.att_neighbor: assert FLAGS.att_type == "my" # These numbers are the target vocabulary sizes of the datasets. # It allows for using different vocabularies for source and targets, # following the implementation in Open-NMT. # I will later put these into command line arguments. if FLAGS.use_bpe: if FLAGS.dataset == "nyt": output_size = 10013 elif FLAGS.dataset == "giga": output_size = 24654 elif FLAGS.dataset == "cnnd": output_size = 10232 else: if FLAGS.dataset == "nyt": output_size = 68885 elif FLAGS.dataset == "giga": output_size = 107389 elif FLAGS.dataset == "cnnd": output_size = 21000 vocab = data.Vocab(FLAGS.vocab_path, FLAGS.vocab_size, FLAGS.dataset) hps = tf.contrib.training.HParams( sample_neighbor=FLAGS.sample_neighbor, use_cluster=FLAGS.use_cluster, binary_neighbor=FLAGS.binary_neighbor, att_neighbor=FLAGS.att_neighbor, encode_neighbor=FLAGS.encode_neighbor, sum_neighbor=FLAGS.sum_neighbor, dataset=FLAGS.dataset, rnn_cell=FLAGS.rnn_cell, output_size=output_size + vocab.offset, train_path=FLAGS.train_path, dev_path=FLAGS.dev_path, tie_embedding=FLAGS.tie_embedding, use_bpe=FLAGS.use_bpe, use_copy=FLAGS.use_copy, reuse_attention=FLAGS.reuse_attention, use_bridge=FLAGS.use_bridge, use_residual=FLAGS.use_residual, att_type=FLAGS.att_type, random_neighbor=FLAGS.random_neighbor, num_neighbors=FLAGS.num_neighbors, model=FLAGS.model, trainer=FLAGS.trainer, learning_rate=FLAGS.learning_rate, lr_schedule=FLAGS.lr_schedule, total_steps=FLAGS.total_steps, emb_dim=FLAGS.emb_dim, binary_dim=FLAGS.binary_dim, neighbor_dim=FLAGS.neighbor_dim, drop=FLAGS.drop, emb_drop=FLAGS.emb_drop, out_drop=FLAGS.out_drop, encoder_drop=FLAGS.encoder_drop, decoder_drop=FLAGS.decoder_drop, weight_decay=FLAGS.weight_decay, encoder_dim=FLAGS.encoder_dim, num_encoder_layers=FLAGS.num_encoder_layers, decoder_dim=FLAGS.decoder_dim, num_decoder_layers=FLAGS.num_decoder_layers, num_mlp_layers=FLAGS.num_mlp_layers, rank=FLAGS.rank, sigma_norm=FLAGS.sigma_norm, batch_size=FLAGS.batch_size, sampling_probability=FLAGS.sampling_probability, beam_width=FLAGS.beam_width, max_enc_steps=FLAGS.max_enc_steps, max_dec_steps=FLAGS.max_dec_steps, vocab_size=FLAGS.vocab_size, max_grad_norm=FLAGS.max_grad_norm, length_norm=FLAGS.length_norm, cp=FLAGS.coverage_penalty, predict_mode=FLAGS.predict_mode) run_config = tf_estimator.RunConfig(model_dir=FLAGS.model_dir) vocab = data.Vocab(FLAGS.vocab_path, FLAGS.vocab_size, FLAGS.dataset) eval_input_fn = partial(data.input_function, is_train=False, vocab=vocab, hps=hps) estimator = tf_estimator.Estimator(model_fn=partial( model_function.model_function, vocab=vocab, hps=hps), config=run_config, model_dir=run_config.model_dir) results = estimator.predict(input_fn=eval_input_fn) with tf.gfile.Open("%s/prediction" % FLAGS.model_dir, "w") as fout: for result in results: outputs, _ = result["outputs"], result["lengths"] prediction = data.id2text(outputs, vocab, use_bpe=FLAGS.use_bpe) fout.write(prediction + "\n")
def main(argv): del argv # Unused. if FLAGS.squared_value: is_squared = 'squared' else: is_squared = 'not_squared' if FLAGS.keep_information: info_keep = 'keep' else: info_keep = 'remove' if FLAGS.dataset_name == 'food_101': params = food_101_params elif FLAGS.dataset_name == 'imagenet': params = imagenet_params elif FLAGS.dataset_name == 'birdsnap': params = birdsnap_params else: raise ValueError('Dataset type is not known %s' % (FLAGS.dataset)) if FLAGS.test_small_sample: model_dir = '/tmp/lalala/' else: model_dir = os.path.join(FLAGS.output_dir, FLAGS.dataset_name, FLAGS.transformation, str(FLAGS.threshold), str(params['base_learning_rate']), str(params['weight_decay']), is_squared, info_keep) if FLAGS.transformation in ['modified_image', 'raw_saliency_map']: model_dir = os.path.join(model_dir, FLAGS.saliency_method) if FLAGS.mode == 'eval': split = 'validation' else: split = 'training' mean_stats = [0.485, 0.456, 0.406] std_stats = [0.229, 0.224, 0.225] update_params = { 'mean_rgb': mean_stats, 'stddev_rgb': std_stats, 'lr_schedule': [ # (multiplier, epoch to start) tuples (1.0, 5), (0.1, 30), (0.01, 60), (0.001, 80) ], 'momentum': 0.9, 'data_format': 'channels_last' } params.update(update_params) sal_method = saliency_dict[FLAGS.saliency_method] if FLAGS.test_small_sample: update_params = { 'train_batch_size': 2, 'eval_batch_size': 2, 'num_train_steps': 10, 'num_images': 2 } params.update(update_params) data_directory = os.path.join(FLAGS.base_dir, FLAGS.dataset_name, '2018-12-10', 'resnet_50', sal_method, split + '*') dataset_ = data_input.DataIterator( mode=FLAGS.mode, data_directory=data_directory, saliency_method=FLAGS.saliency_method, transformation=FLAGS.transformation, threshold=FLAGS.threshold, keep_information=FLAGS.keep_information, use_squared_value=FLAGS.squared_value, mean_stats=mean_stats, std_stats=std_stats, test_small_sample=FLAGS.test_small_sample, num_cores=FLAGS.num_cores) params['output_dir'] = model_dir if FLAGS.mode == 'train': params['batch_size'] = params['train_batch_size'] else: params['batch_size'] = params['eval_batch_size'] num_train_steps = params['num_train_steps'] eval_steps = params['num_eval_images'] // params['batch_size'] run_config = tf_estimator.RunConfig( model_dir=model_dir, save_checkpoints_steps=FLAGS.steps_per_checkpoint) classifier = tf_estimator.Estimator(model_fn=resnet_model_fn, model_dir=model_dir, params=params, config=run_config) if FLAGS.mode == 'eval': # Run evaluation when there's a new checkpoint for ckpt in tf2.training.checkpoints_iterator(model_dir): tf.logging.info('Starting to evaluate.') try: classifier.evaluate(input_fn=dataset_.input_fn, steps=eval_steps, checkpoint_path=ckpt) current_step = int(os.path.basename(ckpt).split('-')[1]) if current_step >= num_train_steps: print('Evaluation finished after training step %d' % current_step) break except tf.errors.NotFoundError: tf.logging.info( 'Checkpoint was not found, skipping checkpoint.') else: if FLAGS.mode == 'train': tf.logging.info('start training...') classifier.train(input_fn=dataset_.input_fn, max_steps=num_train_steps) tf.logging.info('finished training.')
def run_experiment(model_fn, train_input_fn, eval_input_fn, exporters=None, params=None, params_fname=None): """Run an experiment using estimators. This is a light wrapper around typical estimator usage to avoid boilerplate code. Please use the following components separately for more complex usages. Args: model_fn: A model function to be passed to the estimator. See https://www.tensorflow.org/api_docs/python/tf/estimator/Estimator#args_1 train_input_fn: An input function to be passed to the estimator that corresponds to the training data. See https://www.tensorflow.org/api_docs/python/tf/estimator/Estimator#train eval_input_fn: An input function to be passed to the estimator that corresponds to the held-out eval data. See https://www.tensorflow.org/api_docs/python/tf/estimator/Estimator#evaluate exporters: (Optional) An tf.estimator.Exporter or a list of them. params: (Optional) A dictionary of parameters that will be accessible by the model_fn and input_fns. The 'batch_size' and 'use_tpu' values will be set automatically. params_fname: (Optional) If specified, `params` will be written to here under `FLAGS.model_dir` in JSON format. """ params = params if params is not None else {} params.setdefault("use_tpu", FLAGS.use_tpu) if FLAGS.model_dir and params_fname: tf.io.gfile.makedirs(FLAGS.model_dir) params_path = os.path.join(FLAGS.model_dir, params_fname) with tf.io.gfile.GFile(params_path, "w") as params_file: json.dump(params, params_file, indent=2, sort_keys=True) if params["use_tpu"]: if FLAGS.tpu_name: tpu_cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver( FLAGS.tpu_name, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project) else: tpu_cluster_resolver = None run_config = tf_estimator.tpu.RunConfig( cluster=tpu_cluster_resolver, master=FLAGS.master, model_dir=FLAGS.model_dir, tf_random_seed=FLAGS.tf_random_seed, save_checkpoints_steps=FLAGS.save_checkpoints_steps, tpu_config=tf_estimator.tpu.TPUConfig( iterations_per_loop=FLAGS.save_checkpoints_steps)) if "batch_size" in params: # Let the TPUEstimator fill in the batch size. params.pop("batch_size") estimator = tf_estimator.tpu.TPUEstimator( use_tpu=True, model_fn=model_fn, params=params, config=run_config, train_batch_size=FLAGS.batch_size, eval_batch_size=FLAGS.eval_batch_size, predict_batch_size=FLAGS.eval_batch_size) else: run_config = tf_estimator.RunConfig( model_dir=FLAGS.model_dir, tf_random_seed=FLAGS.tf_random_seed, save_checkpoints_steps=FLAGS.save_checkpoints_steps, keep_checkpoint_max=FLAGS.keep_checkpoint_max) params["batch_size"] = FLAGS.batch_size estimator = tf_estimator.Estimator( config=run_config, model_fn=model_fn, params=params, model_dir=FLAGS.model_dir) train_spec = tf_estimator.TrainSpec( input_fn=train_input_fn, max_steps=FLAGS.num_train_steps) eval_spec = tf_estimator.EvalSpec( name="default", input_fn=eval_input_fn, exporters=exporters, start_delay_secs=FLAGS.eval_start_delay_secs, throttle_secs=FLAGS.eval_throttle_secs, steps=FLAGS.num_eval_steps) tf.logging.set_verbosity(tf.logging.INFO) tf_estimator.train_and_evaluate( estimator=estimator, train_spec=train_spec, eval_spec=eval_spec)
def run_model(): """Run experiment with tf.estimator. """ params = { 'kb_index': FLAGS.kb_index, 'cm_width': FLAGS.cm_width, 'cm_depth': FLAGS.cm_depth, 'entity_emb_size': FLAGS.entity_emb_size, 'relation_emb_size': FLAGS.relation_emb_size, 'vocab_emb_size': FLAGS.vocab_emb_size, 'max_set': FLAGS.max_set, 'learning_rate': FLAGS.learning_rate, 'gradient_clip': FLAGS.gradient_clip, 'intermediate_top_k': FLAGS.intermediate_top_k, 'use_cm_sketch': FLAGS.use_cm_sketch, 'train_entity_emb': FLAGS.train_entity_emb, 'train_relation_emb': FLAGS.train_relation_emb, 'bert_handle': FLAGS.bert_handle, 'train_bert': FLAGS.train_bert, } data_loader = DataLoader(params, FLAGS.name, get_root_dir(FLAGS.name), FLAGS.kb_file, FLAGS.vocab_file) estimator_config = tf_estimator.RunConfig( save_checkpoints_steps=FLAGS.checkpoint_step) warm_start_settings = tf_estimator.WarmStartSettings( # pylint: disable=g-long-ternary ckpt_to_initialize_from=FLAGS.load_model_dir, vars_to_warm_start=[ 'embeddings_mat/entity_embeddings_mat', 'embeddings_mat/relation_embeddings_mat' ], ) if FLAGS.load_model_dir is not None else None estimator = tf_estimator.Estimator( model_fn=build_model_fn(FLAGS.name, data_loader, FLAGS.eval_name, FLAGS.eval_metric_at_k), model_dir=FLAGS.checkpoint_dir + FLAGS.model_name, config=estimator_config, params=params, warm_start_from=warm_start_settings) if FLAGS.mode == 'train': train_input_fn = data_loader.build_input_fn( name=FLAGS.name, batch_size=FLAGS.batch_size, mode='train', epochs=FLAGS.epochs, n_take=-1, shuffle=True) eval_input_fn = data_loader.build_input_fn(name=FLAGS.name, batch_size=FLAGS.batch_size, mode='eval', epochs=1, n_take=FLAGS.num_eval, shuffle=False) # Define mode-specific operations if FLAGS.mode == 'train': train_spec = tf_estimator.TrainSpec(input_fn=train_input_fn) # Busy waiting for evaluation until new checkpoint comes out test_spec = tf_estimator.EvalSpec(input_fn=eval_input_fn, steps=FLAGS.num_online_eval, start_delay_secs=0, throttle_secs=FLAGS.eval_time) tf_estimator.train_and_evaluate(estimator, train_spec, test_spec) elif FLAGS.mode == 'eval': tf_evaluation = estimator.evaluate(eval_input_fn) print(tf_evaluation) elif FLAGS.mode == 'pred': tf_predictions = estimator.predict(eval_input_fn) if FLAGS.name.startswith('query2box'): task = FLAGS.name.split('_')[-1] metrics = Query2BoxMetrics(task, FLAGS.root_dir, data_loader) else: raise NotImplementedError() for tf_prediction in tqdm(tf_predictions): metrics.eval(tf_prediction) metrics.print_metrics() else: raise ValueError('mode not recognized: %s' % FLAGS.mode)
def main(_): # Modify the paths to save results when tuning hyperparameters. if FLAGS.node_encoder == "lstm": FLAGS.result_path = os.path.join(FLAGS.result_path, str(FLAGS.node_lstm_size)) if FLAGS.node_encoder == "transformer": FLAGS.result_path = os.path.join( FLAGS.result_path, "max_steps_" + str(FLAGS.max_steps_no_increase)) FLAGS.result_path = os.path.join( FLAGS.result_path, "hidden_unit_" + str(FLAGS.transformer_hidden_unit)) if FLAGS.cross_vertical: FLAGS.result_path = os.path.join( FLAGS.result_path, "CKP-{0}/{1}/".format(FLAGS.checkpoint_vertical, FLAGS.checkpoint_websites)) FLAGS.checkpoint_path = os.path.join( FLAGS.checkpoint_path, "{0}/{1}-results/".format(FLAGS.checkpoint_vertical, FLAGS.checkpoint_websites)) tf.gfile.MakeDirs( os.path.join( FLAGS.result_path, "{0}/{1}-results/".format(FLAGS.vertical, FLAGS.source_website))) tf.logging.set_verbosity(tf.logging.INFO) if FLAGS.use_uniform_embedding: vocab_vertical = "all" else: vocab_vertical = FLAGS.vertical # Hyper-parameters. params = { "add_goldmine": FLAGS.add_goldmine, "add_leaf_types": FLAGS.add_leaf_types, "batch_size": FLAGS.batch_size, "buffer": 1000, # Buffer for shuffling. No need to care about. "chars": os.path.join(FLAGS.domtree_data_path, "%s.vocab.chars.txt" % vocab_vertical), "circle_features": FLAGS.circle_features, "dim_word_embedding": FLAGS.dim_word_embedding, "dim_chars": FLAGS.dim_chars, "dim_label_embedding": FLAGS.dim_label_embedding, "dim_goldmine": 30, "dim_leaf_type": 20, "dim_positions": 30, "dim_xpath_units": FLAGS.dim_xpath_units, "dropout": 0.3, "epochs": FLAGS.epochs, "extract_node_emb": FLAGS.extract_node_emb, "filters": 50, # The dimension of char-level word representations. "friend_encoder": FLAGS.friend_encoder, "use_friend_semantic": FLAGS.use_friend_semantic, "goldmine_features": os.path.join(FLAGS.domtree_data_path, "vocab.goldmine_features.txt"), "glove": os.path.join( FLAGS.domtree_data_path, "%s.%d.emb.npz" % (vocab_vertical, FLAGS.dim_word_embedding)), "friend_hidden_size": FLAGS.friend_hidden_size, "kernel_size": 3, # CNN window size to embed char sequences. "last_hidden_layer_size": FLAGS.last_hidden_layer_size, "leaf_types": os.path.join(FLAGS.domtree_data_path, "%s.vocab.leaf_types.txt" % vocab_vertical), "lstm_size": 100, "max_steps_no_increase": FLAGS.max_steps_no_increase, "node_encoder": FLAGS.node_encoder, "node_filters": 100, "node_kernel_size": 5, "node_lstm_size": FLAGS.node_lstm_size, "num_oov_buckets": 1, "objective": FLAGS.objective, "positions": os.path.join(FLAGS.domtree_data_path, "vocab.positions.txt"), "running_mode": FLAGS.run, "semantic_encoder": FLAGS.semantic_encoder, "source_website": FLAGS.source_website, "tags": os.path.join(FLAGS.domtree_data_path, "%s.vocab.tags.txt" % (FLAGS.vertical)), "tags-all": os.path.join(FLAGS.domtree_data_path, "all.vocab.tags.txt"), "target_website": FLAGS.target_website, "transformer_hidden_unit": FLAGS.transformer_hidden_unit, "transformer_head": FLAGS.transformer_head, "transformer_hidden_layer": FLAGS.transformer_hidden_layer, "use_crf": FLAGS.use_crf, "use_friends_cnn": FLAGS.use_friends_cnn, "use_friends_discrete_feature": FLAGS.use_friends_discrete_feature, "use_prev_text_lstm": FLAGS.use_prev_text_lstm, "use_xpath_lstm": FLAGS.use_xpath_lstm, "use_uniform_label": FLAGS.use_uniform_label, "use_position_embedding": FLAGS.use_position_embedding, "words": os.path.join(FLAGS.domtree_data_path, "%s.vocab.words.txt" % vocab_vertical), "xpath_lstm_size": 100, "xpath_units": os.path.join(FLAGS.domtree_data_path, "%s.vocab.xpath_units.txt" % vocab_vertical), } with tf.gfile.Open( os.path.join( FLAGS.result_path, "{0}/{1}-results/params.json".format(FLAGS.vertical, FLAGS.source_website)), "w") as f: json.dump(params, f, indent=4, sort_keys=True) # Build estimator, train and evaluate. train_input_function = functools.partial( model_util.joint_input_fn, get_data_path(vertical=FLAGS.vertical, website=FLAGS.source_website, dev=False, goldmine=False), get_data_path(vertical=FLAGS.vertical, website=FLAGS.source_website, dev=False, goldmine=True), FLAGS.vertical, params, shuffle_and_repeat=True, mode="train") cfg = tf_estimator.RunConfig(save_checkpoints_steps=300, save_summary_steps=300, tf_random_seed=42) # Set up the checkpoint to load. if FLAGS.checkpoint_path: # The best model was always saved in "cpkt-601". checkpoint_file = FLAGS.checkpoint_path + "/model/model.ckpt-601" # Do not load parameters whose names contain the "label_dense". # These parameters are ought to be learned from scratch. ws = tf_estimator.WarmStartSettings( ckpt_to_initialize_from=checkpoint_file, vars_to_warm_start="^((?!label_dense).)*$") estimator = tf_estimator.Estimator(models.joint_extraction_model_fn, os.path.join( FLAGS.result_path, "{0}/{1}-results/model".format( FLAGS.vertical, FLAGS.source_website)), cfg, params, warm_start_from=ws) else: estimator = tf_estimator.Estimator( models.joint_extraction_model_fn, os.path.join( FLAGS.result_path, "{0}/{1}-results/model".format(FLAGS.vertical, FLAGS.source_website)), cfg, params) tf.gfile.MakeDirs(estimator.eval_dir()) hook = early_stopping.stop_if_no_increase_hook( estimator, metric_name="f1", max_steps_without_increase=FLAGS.max_steps_no_increase, min_steps=300, run_every_steps=100, run_every_secs=None) train_spec = tf_estimator.TrainSpec(input_fn=train_input_function, hooks=[hook]) if FLAGS.run == "train": eval_input_function = functools.partial( model_util.joint_input_fn, get_data_path(vertical=FLAGS.vertical, website=FLAGS.source_website, dev=True, goldmine=False), get_data_path(vertical=FLAGS.vertical, website=FLAGS.source_website, dev=True, goldmine=True), FLAGS.vertical, mode="all") eval_spec = tf_estimator.EvalSpec(input_fn=eval_input_function, steps=300, throttle_secs=1) tf_estimator.train_and_evaluate(estimator, train_spec, eval_spec) target_websites = FLAGS.target_website.split("_") if FLAGS.source_website not in target_websites: target_websites = [FLAGS.source_website] + target_websites for target_website in target_websites: write_predictions(estimator=estimator, vertical=FLAGS.vertical, source_website=FLAGS.source_website, target_website=target_website) model_util.page_hits_level_metric(result_path=FLAGS.result_path, vertical=FLAGS.vertical, source_website=FLAGS.source_website, target_website=target_website) model_util.site_level_voting(result_path=FLAGS.result_path, vertical=FLAGS.vertical, source_website=FLAGS.source_website, target_website=target_website) model_util.page_level_constraint( domtree_data_path=FLAGS.domtree_data_path, result_path=FLAGS.result_path, vertical=FLAGS.vertical, source_website=FLAGS.source_website, target_website=target_website)