def __init__(self, head, subnetwork_generator, max_iteration_steps, ensemblers=None, ensemble_strategies=None, evaluator=None, report_materializer=None, metric_fn=None, force_grow=False, replicate_ensemble_in_training=False, adanet_loss_decay=.9, model_dir=None, report_dir=None, config=None, use_tpu=True, train_batch_size=None, eval_batch_size=None, debug=False, **kwargs): if tf_compat.version_greater_or_equal("2.0.0"): raise ValueError( "TPUEstimator is not yet supported with TensorFlow 2.0.") self._use_tpu = use_tpu if not self._use_tpu: tf.logging.warning( "This adanet.TPUEstimator is meant to be used for running on TPU. " "If you want to run on CPU/GPU, use adanet.Estimator instead.") # TPUEstimator modifies config under the hood. We keep track of it here so # we can use it during the bookkeeping phase and when predict() is called. self._original_config = config or tf.contrib.RunConfig() self._train_batch_size = train_batch_size or 0 self._eval_batch_size = eval_batch_size or train_batch_size or 0 super(TPUEstimator, self).__init__( head=head, subnetwork_generator=subnetwork_generator, max_iteration_steps=max_iteration_steps, ensemblers=ensemblers, ensemble_strategies=ensemble_strategies, evaluator=evaluator, report_materializer=report_materializer, metric_fn=metric_fn, force_grow=force_grow, replicate_ensemble_in_training=replicate_ensemble_in_training, adanet_loss_decay=adanet_loss_decay, model_dir=model_dir, report_dir=report_dir, config=self._original_config, use_tpu=use_tpu, eval_on_tpu=use_tpu, export_to_tpu=False, train_batch_size=self._train_batch_size, eval_batch_size=self._eval_batch_size, **kwargs)
def create_eval_metrics(self, features, labels, estimator_spec, metric_fn): """Creates evaluation metrics from the given arguments. Args: features: Input `dict` of `Tensor` objects. labels: Labels `Tensor` or a dictionary of string label name to `Tensor` (for multi-head). estimator_spec: The `EstimatorSpec` created by a `Head` instance. metric_fn: A function which should obey the following signature: - Args: can only have following three arguments in any order: * predictions: Predictions `Tensor` or dict of `Tensor` created by given `Head`. * features: Input `dict` of `Tensor` objects created by `input_fn` which is given to `estimator.evaluate` as an argument. * labels: Labels `Tensor` or dict of `Tensor` (for multi-head) created by `input_fn` which is given to `estimator.evaluate` as an argument. - Returns: Dict of metric results keyed by name. Final metrics are a union of this and `estimator`s existing metrics. If there is a name conflict between this and `estimator`s existing metrics, this will override the existing one. The values of the dict are the results of calling a metric function, namely a `(metric_tensor, update_op)` tuple. """ # If estimator_spec is not a TPUEstimatorSpec we create dummy metric_fn # and args. if isinstance(estimator_spec, tf.estimator.EstimatorSpec): spec_fn, spec_args = lambda: estimator_spec.eval_metric_ops, [] else: spec_fn, spec_args = estimator_spec.eval_metrics self._eval_metrics_store.add_eval_metrics( self._templatize_metric_fn(spec_fn), spec_args) if tf_compat.version_greater_or_equal( "1.13.0") and tf.executing_eagerly(): loss_metric = tf.keras.metrics.Mean("mean_loss") def loss_fn(loss): loss_metric(loss) return {"loss": loss_metric} else: loss_fn = lambda loss: {"loss": tf_compat.v1.metrics.mean(loss)} loss_fn_args = [tf.reshape(estimator_spec.loss, [1])] self._eval_metrics_store.add_eval_metrics( self._templatize_metric_fn(loss_fn), loss_fn_args) # NOTE: the user supplied metrics_fn must be added last. This is because we # want user metrics to override AdaNet's metrics. if metric_fn: metric_fn_args = {} argspec = inspect.getargspec(metric_fn).args if "features" in argspec: metric_fn_args["features"] = features if "labels" in argspec: metric_fn_args["labels"] = labels if "predictions" in argspec: metric_fn_args["predictions"] = estimator_spec.predictions self._eval_metrics_store.add_eval_metrics( self._templatize_metric_fn(metric_fn), metric_fn_args)
def setUp(self): super(TPUEstimatorTest, self).setUp() if not tf_compat.version_greater_or_equal("1.14.0"): self.skipTest("TPUEmbedding not supported in version 1.13.0 and below.") # TPUConfig initializes model_dir from TF_CONFIG and checks that the user # provided model_dir matches the TF_CONFIG one. tf_config = {"model_dir": self.test_subdirectory} os.environ["TF_CONFIG"] = json.dumps(tf_config)
def monkey_patched_summaries(summary): """A context where global summary functions point to the given summary. Restores original summary functions upon exit. NOTE: This function is not thread-safe. Args: summary: An `adanet.Summary` instance. Yields: A context where summary functions are routed to the given `adanet.Summary`. """ from tensorflow.python.ops import summary_ops_v2 as summary_v2_lib # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top old_summary_scalar = summary_lib.scalar old_summary_image = summary_lib.image old_summary_histogram = summary_lib.histogram old_summary_audio = summary_lib.audio old_summary_v2_scalar = summary_v2_lib.scalar old_summary_v2_image = summary_v2_lib.image old_summary_v2_histogram = summary_v2_lib.histogram old_summary_v2_audio = summary_v2_lib.audio old_summary_compat_v2_scalar = tf_compat.v2.summary.scalar old_summary_compat_v2_image = tf_compat.v2.summary.image old_summary_compat_v2_histogram = tf_compat.v2.summary.histogram old_summary_compat_v2_audio = tf_compat.v2.summary.audio # Monkey-patch global attributes. wrapped_summary = _SummaryWrapper(summary) setattr(tf_v1.summary, "scalar", wrapped_summary.scalar) setattr(tf_v1.summary, "image", wrapped_summary.image) setattr(tf_v1.summary, "histogram", wrapped_summary.histogram) setattr(tf_v1.summary, "audio", wrapped_summary.audio) setattr(tf_compat.v1.summary, "scalar", wrapped_summary.scalar) setattr(tf_compat.v1.summary, "image", wrapped_summary.image) setattr(tf_compat.v1.summary, "histogram", wrapped_summary.histogram) setattr(tf_compat.v1.summary, "audio", wrapped_summary.audio) setattr(summary_lib, "scalar", wrapped_summary.scalar) setattr(summary_lib, "image", wrapped_summary.image) setattr(summary_lib, "histogram", wrapped_summary.histogram) setattr(summary_lib, "audio", wrapped_summary.audio) setattr(tf_compat.v2.summary, "scalar", wrapped_summary.scalar_v3) setattr(tf_compat.v2.summary, "image", wrapped_summary.image_v3) setattr(tf_compat.v2.summary, "histogram", wrapped_summary.histogram_v3) setattr(tf_compat.v2.summary, "audio", wrapped_summary.audio_v3) setattr(summary_v2_lib, "scalar", wrapped_summary.scalar_v2) setattr(summary_v2_lib, "image", wrapped_summary.image_v2) setattr(summary_v2_lib, "histogram", wrapped_summary.histogram_v2) setattr(summary_v2_lib, "audio", wrapped_summary.audio_v2) try: # TF 2.0 eliminates tf.contrib. setattr(tf_v1.contrib.summary, "scalar", wrapped_summary.scalar_v2) setattr(tf_v1.contrib.summary, "image", wrapped_summary.image_v2) setattr(tf_v1.contrib.summary, "histogram", wrapped_summary.histogram_v2) setattr(tf_v1.contrib.summary, "audio", wrapped_summary.audio_v2) except (AttributeError, ImportError): # TF 2.0 eliminates tf.contrib. # Also set the new tf.summary to be use the new summaries in TF 2. if tf_compat.version_greater_or_equal("2.0.0"): setattr(tf.summary, "scalar", wrapped_summary.scalar_v3) setattr(tf.summary, "image", wrapped_summary.image_v3) setattr(tf.summary, "histogram", wrapped_summary.histogram_v3) setattr(tf.summary, "audio", wrapped_summary.audio_v3) try: yield finally: # Revert monkey-patches. try: setattr(tf_v1.contrib.summary, "audio", old_summary_v2_audio) setattr(tf_v1.contrib.summary, "histogram", old_summary_v2_histogram) setattr(tf_v1.contrib.summary, "image", old_summary_v2_image) setattr(tf_v1.contrib.summary, "scalar", old_summary_v2_scalar) except (AttributeError, ImportError): # TF 2.0 eliminates tf.contrib. pass setattr(summary_v2_lib, "audio", old_summary_v2_audio) setattr(summary_v2_lib, "histogram", old_summary_v2_histogram) setattr(summary_v2_lib, "image", old_summary_v2_image) setattr(summary_v2_lib, "scalar", old_summary_v2_scalar) setattr(tf.summary, "audio", old_summary_compat_v2_audio) setattr(tf.summary, "histogram", old_summary_compat_v2_histogram) setattr(tf.summary, "image", old_summary_compat_v2_image) setattr(tf.summary, "scalar", old_summary_compat_v2_scalar) setattr(tf_compat.v2.summary, "audio", old_summary_compat_v2_audio) setattr(tf_compat.v2.summary, "histogram", old_summary_compat_v2_histogram) setattr(tf_compat.v2.summary, "image", old_summary_compat_v2_image) setattr(tf_compat.v2.summary, "scalar", old_summary_compat_v2_scalar) setattr(summary_lib, "audio", old_summary_audio) setattr(summary_lib, "histogram", old_summary_histogram) setattr(summary_lib, "image", old_summary_image) setattr(summary_lib, "scalar", old_summary_scalar) setattr(tf_compat.v1.summary, "audio", old_summary_audio) setattr(tf_compat.v1.summary, "histogram", old_summary_histogram) setattr(tf_compat.v1.summary, "image", old_summary_image) setattr(tf_compat.v1.summary, "scalar", old_summary_scalar) setattr(tf_v1.summary, "audio", old_summary_audio) setattr(tf_v1.summary, "histogram", old_summary_histogram) setattr(tf_v1.summary, "image", old_summary_image) setattr(tf_v1.summary, "scalar", old_summary_scalar)
def test_auto_ensemble_estimator_lifecycle(self, list_candidate_pool): features = {"input_1": [[1., 0.]]} labels = [[1.]] run_config = tf.estimator.RunConfig(tf_random_seed=42) head = tf.contrib.estimator.regression_head( loss_reduction=tf.losses.Reduction.SUM_OVER_BATCH_SIZE) optimizer = tf.train.GradientDescentOptimizer(learning_rate=.01) feature_columns = [tf.feature_column.numeric_column("input_1", shape=[2])] def train_input_fn(): input_features = {} for key, feature in features.items(): input_features[key] = tf.constant(feature, name=key) input_labels = tf.constant(labels, name="labels") return input_features, input_labels def test_input_fn(): input_features = tf.data.Dataset.from_tensors([ tf.constant(features["input_1"]) ]).make_one_shot_iterator().get_next() return {"input_1": input_features}, None if hasattr(tf.estimator, "LinearEstimator"): linear_estimator_fn = tf.estimator.LinearEstimator else: linear_estimator_fn = tf.contrib.estimator.LinearEstimator if hasattr(tf.estimator, "DNNEstimator"): dnn_estimator_fn = tf.estimator.DNNEstimator else: dnn_estimator_fn = tf.contrib.estimator.DNNEstimator candidate_pool = { "linear": linear_estimator_fn( head=head, feature_columns=feature_columns, optimizer=optimizer), "dnn": dnn_estimator_fn( head=head, feature_columns=feature_columns, optimizer=optimizer, hidden_units=[3]) } if list_candidate_pool: candidate_pool = [candidate_pool[k] for k in sorted(candidate_pool)] estimator = AutoEnsembleEstimator( head=head, candidate_pool=candidate_pool, max_iteration_steps=10, force_grow=True, model_dir=self.test_subdirectory, config=run_config) # Train for three iterations. estimator.train(input_fn=train_input_fn, max_steps=30) # Evaluate. eval_results = estimator.evaluate(input_fn=train_input_fn, steps=3) want_loss = .209 if tf_compat.version_greater_or_equal("1.10.0") and ( not tf_compat.version_greater_or_equal("1.12.0")): # Only TF 1.10 and 1.11. want_loss = .079514 self.assertAllClose(want_loss, eval_results["loss"], atol=.05) # Predict. predictions = estimator.predict(input_fn=test_input_fn) for prediction in predictions: self.assertIsNotNone(prediction["predictions"]) # Export SavedModel. def serving_input_fn(): """Input fn for serving export, starting from serialized example.""" serialized_example = tf.placeholder( dtype=tf.string, shape=(None), name="serialized_example") for key, value in features.items(): features[key] = tf.constant(value) return export.SupervisedInputReceiver( features=features, labels=tf.constant(labels), receiver_tensors=serialized_example) export_dir_base = os.path.join(self.test_subdirectory, "export") export_saved_model_fn = getattr(estimator, "export_saved_model", None) if not callable(export_saved_model_fn): export_saved_model_fn = estimator.export_savedmodel export_saved_model_fn( export_dir_base=export_dir_base, serving_input_receiver_fn=serving_input_fn)
def train_and_evaluate_estimator(): """Runs Estimator distributed training.""" # The tf.estimator.RunConfig automatically parses the TF_CONFIG environment # variables during construction. # For more information on how tf.estimator.RunConfig uses TF_CONFIG, see # https://www.tensorflow.org/api_docs/python/tf/estimator/RunConfig. config = tf.estimator.RunConfig( tf_random_seed=42, model_dir=FLAGS.model_dir, session_config=tf_compat.v1.ConfigProto( log_device_placement=False, # Ignore other workers; only talk to parameter servers. # Otherwise, when a chief/worker terminates, the others will hang. device_filters=["/job:ps"])) head = regression_head.RegressionHead( loss_reduction=tf_compat.v2.losses.Reduction.SUM_OVER_BATCH_SIZE) kwargs = { "max_iteration_steps": 100, "force_grow": True, "delay_secs_per_worker": .2, "max_worker_delay_secs": 1, "worker_wait_secs": .5, # Set low timeout to reduce wait time for failures. "worker_wait_timeout_secs": 60, "config": config } if FLAGS.placement_strategy == "round_robin": kwargs["experimental_placement_strategy"] = RoundRobinStrategy() if FLAGS.estimator_type == "autoensemble": feature_columns = [tf.feature_column.numeric_column("x", shape=[2])] if hasattr(tf.estimator, "LinearEstimator"): linear_estimator_fn = tf_compat.v1.estimator.LinearEstimator else: linear_estimator_fn = tf.contrib.estimator.LinearEstimator if hasattr(tf.estimator, "DNNEstimator"): dnn_estimator_fn = tf_compat.v1.estimator.DNNEstimator else: dnn_estimator_fn = tf.contrib.estimator.DNNEstimator candidate_pool = { "linear": linear_estimator_fn(head=head, feature_columns=feature_columns, optimizer=tf_compat.v1.train.AdamOptimizer( learning_rate=.001)), "dnn": dnn_estimator_fn( head=head, feature_columns=feature_columns, optimizer=tf_compat.v1.train.AdamOptimizer(learning_rate=.001), hidden_units=[3]), "dnn2": dnn_estimator_fn( head=head, feature_columns=feature_columns, optimizer=tf_compat.v1.train.AdamOptimizer(learning_rate=.001), hidden_units=[5]), } estimator = AutoEnsembleEstimator(head=head, candidate_pool=candidate_pool, **kwargs) elif FLAGS.estimator_type == "estimator": subnetwork_generator = SimpleGenerator([ _DNNBuilder("dnn1", config, layer_size=3), _DNNBuilder("dnn2", config, layer_size=4), _DNNBuilder("dnn3", config, layer_size=5), ]) estimator = Estimator(head=head, subnetwork_generator=subnetwork_generator, **kwargs) def input_fn(): xor_features = [[1., 0.], [0., 0], [0., 1.], [1., 1.]] xor_labels = [[1.], [0.], [1.], [0.]] input_features = {"x": tf.constant(xor_features, name="x")} input_labels = tf.constant(xor_labels, name="y") return input_features, input_labels train_hooks = [] # ProfilerHook raises the following error in older TensorFlow versions: # ValueError: The provided tag was already used for this event type. if tf_compat.version_greater_or_equal("1.13.0"): train_hooks = [ tf.estimator.ProfilerHook(save_steps=50, output_dir=FLAGS.model_dir) ] # Train for three iterations. train_spec = tf.estimator.TrainSpec(input_fn=input_fn, max_steps=300, hooks=train_hooks) eval_spec = tf.estimator.EvalSpec(input_fn=input_fn, steps=1, start_delay_secs=.5, throttle_secs=.5) # Calling train_and_evaluate is the official way to perform distributed # training with an Estimator. Calling Estimator#train directly results # in an error when the TF_CONFIG is setup for a cluster. tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)