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)
