def test_extra_checkpoint_saver_hook(self):
"""Tests b/122795064."""
features = {"input_1": [[1., 0.]]}
labels = [[1.]]
run_config = tf.estimator.RunConfig(tf_random_seed=42)
head = tf.contrib.estimator.binary_classification_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])
]
estimator = AutoEnsembleEstimator(
head=head,
candidate_pool=[
tf.estimator.LinearClassifier(n_classes=2,
feature_columns=feature_columns,
optimizer=optimizer),
tf.estimator.DNNClassifier(n_classes=2,
feature_columns=feature_columns,
optimizer=optimizer,
hidden_units=[3]),
],
max_iteration_steps=3,
force_grow=True,
model_dir=self.test_subdirectory,
config=run_config)
ckpt_dir = os.path.join(self.test_subdirectory)
hooks = [tf.train.CheckpointSaverHook(ckpt_dir, save_steps=1)]
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
estimator.train(input_fn=train_input_fn, max_steps=6, hooks=hooks)
def test_auto_ensemble_estimator_lifecycle(self,
candidate_pool,
want_loss,
max_train_steps=30):
features = {"input_1": [[1., 0.]]}
labels = [[1.]]
run_config = tf.estimator.RunConfig(tf_random_seed=42)
head = regression_head.RegressionHead()
# Always create optimizers in a lambda to prevent error like:
# `RuntimeError: Cannot set `iterations` to a new Variable after the
# Optimizer weights have been created`
optimizer = lambda: tf.keras.optimizers.SGD(lr=.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():
dataset = tf.data.Dataset.from_tensors([tf.constant(features["input_1"])])
input_features = tf.compat.v1.data.make_one_shot_iterator(
dataset).get_next()
return {"input_1": input_features}, None
estimator = AutoEnsembleEstimator(
head=head,
candidate_pool=candidate_pool(head, feature_columns, optimizer),
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=max_train_steps)
# Evaluate.
eval_results = estimator.evaluate(input_fn=train_input_fn, steps=1)
self.assertAllClose(max_train_steps, eval_results["global_step"])
self.assertAllClose(want_loss, eval_results["loss"], atol=.3)
# 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.compat.v1.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")
estimator.export_saved_model(
export_dir_base=export_dir_base,
serving_input_receiver_fn=serving_input_fn)
def test_auto_ensemble_estimator_lifecycle(self,
candidate_pool,
want_loss,
max_train_steps=30):
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
estimator = AutoEnsembleEstimator(head=head,
candidate_pool=candidate_pool(
head, feature_columns,
optimizer),
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=max_train_steps)
# Evaluate.
eval_results = estimator.evaluate(input_fn=train_input_fn, steps=1)
self.assertAllClose(max_train_steps, eval_results["global_step"])
self.assertAllClose(want_loss, eval_results["loss"], atol=.3)
# 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,
save_checkpoints_steps=10,
save_checkpoints_secs=None,
# Keep all checkpoints to avoid checkpoint GC causing failures during
# evaluation.
# TODO: Prevent checkpoints that are currently being
# evaluated by another process from being garbage collected.
keep_checkpoint_max=None,
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"]))
def input_fn():
input_features = {"x": tf.constant(features, name="x")}
input_labels = tf.constant(labels, name="y")
return tf.data.Dataset.from_tensors((input_features, input_labels)).repeat()
kwargs = {
"max_iteration_steps": 100,
"force_grow": True,
"delay_secs_per_worker": .2,
"max_worker_delay_secs": 1,
"worker_wait_secs": 1,
# Set low timeout to reduce wait time for failures.
"worker_wait_timeout_secs": 180,
"evaluator": Evaluator(input_fn, steps=10),
"config": config
}
head = head_lib._regression_head( # pylint: disable=protected-access
loss_reduction=tf.losses.Reduction.SUM_OVER_BATCH_SIZE)
features = [[1., 0.], [0., 0], [0., 1.], [1., 1.]]
labels = [[1.], [0.], [1.], [0.]]
estimator_type = FLAGS.estimator_type
if FLAGS.placement_strategy == "round_robin":
kwargs["experimental_placement_strategy"] = RoundRobinStrategy()
if estimator_type == "autoensemble":
feature_columns = [tf.feature_column.numeric_column("x", shape=[2])]
# pylint: disable=g-long-lambda
# TODO: Switch optimizers to tf.keras.optimizers.Adam once the
# distribution bug is fixed.
candidate_pool = {
"linear":
tf.estimator.LinearEstimator(
head=head,
feature_columns=feature_columns,
optimizer=lambda: tf_compat.v1.train.AdamOptimizer(
learning_rate=.001)),
"dnn":
tf.estimator.DNNEstimator(
head=head,
feature_columns=feature_columns,
optimizer=lambda: tf_compat.v1.train.AdamOptimizer(
learning_rate=.001),
hidden_units=[3]),
"dnn2":
tf.estimator.DNNEstimator(
head=head,
feature_columns=feature_columns,
optimizer=lambda: tf_compat.v1.train.AdamOptimizer(
learning_rate=.001),
hidden_units=[10, 10]),
}
# pylint: enable=g-long-lambda
estimator = AutoEnsembleEstimator(
head=head, candidate_pool=candidate_pool, **kwargs)
elif 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)
elif FLAGS.estimator_type == "autoensemble_trees_multiclass":
if not bt_losses:
logging.warning(
"Skipped autoensemble_trees_multiclass test since contrib is missing."
)
return
n_classes = 3
head = head_lib._multi_class_head_with_softmax_cross_entropy_loss( # pylint: disable=protected-access
n_classes=n_classes,
loss_reduction=tf.losses.Reduction.SUM_OVER_BATCH_SIZE)
def tree_loss_fn(labels, logits):
result = bt_losses.per_example_maxent_loss(
labels=labels, logits=logits, num_classes=n_classes, weights=None)
return result[0]
tree_head = head_lib._multi_class_head_with_softmax_cross_entropy_loss( # pylint: disable=protected-access
loss_fn=tree_loss_fn,
n_classes=n_classes,
loss_reduction=tf.losses.Reduction.SUM_OVER_BATCH_SIZE)
labels = [[1], [0], [1], [2]]
feature_columns = [tf.feature_column.numeric_column("x", shape=[2])]
# TODO: Switch optimizers to tf.keras.optimizers.Adam once the
# distribution bug is fixed.
candidate_pool = lambda config: { # pylint: disable=g-long-lambda
"linear":
tf.estimator.LinearEstimator(
head=head,
feature_columns=feature_columns,
optimizer=tf_compat.v1.train.AdamOptimizer(
learning_rate=.001),
config=config),
"gbdt":
tf.estimator.BoostedTreesEstimator(
head=tree_head,
feature_columns=feature_columns,
n_trees=10,
n_batches_per_layer=1,
center_bias=False,
config=config),
}
estimator = AutoEnsembleEstimator(
head=head, candidate_pool=candidate_pool, **kwargs)
elif estimator_type == "estimator_with_experimental_multiworker_strategy":
def _model_fn(features, labels, mode):
"""Test model_fn."""
layer = tf.keras.layers.Dense(1)
logits = layer(features["x"])
if mode == tf.estimator.ModeKeys.PREDICT:
predictions = {"logits": logits}
return tf.estimator.EstimatorSpec(mode, predictions=predictions)
loss = tf.losses.mean_squared_error(
labels=labels,
predictions=logits,
reduction=tf.losses.Reduction.SUM_OVER_BATCH_SIZE)
if mode == tf.estimator.ModeKeys.EVAL:
return tf.estimator.EstimatorSpec(mode, loss=loss)
if mode == tf.estimator.ModeKeys.TRAIN:
optimizer = tf.train.GradientDescentOptimizer(0.2)
train_op = optimizer.minimize(
loss, global_step=tf.train.get_global_step())
return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)
if json.loads(os.environ["TF_CONFIG"])["task"]["type"] == "evaluator":
# The evaluator job would crash if MultiWorkerMirroredStrategy is called.
distribution = None
else:
distribution = tf.distribute.experimental.MultiWorkerMirroredStrategy()
multiworker_config = tf.estimator.RunConfig(
tf_random_seed=42,
model_dir=FLAGS.model_dir,
train_distribute=distribution,
session_config=tf_compat.v1.ConfigProto(log_device_placement=False))
# TODO: Replace with adanet.Estimator. Currently this just verifies
# that the distributed testing framework supports distribute strategies.
estimator = tf.estimator.Estimator(
model_fn=_model_fn, config=multiworker_config)
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=.05)
# 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)
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.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 = tf.contrib.estimator.regression_head(
loss_reduction=tf.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.estimator_type == "autoensemble":
feature_columns = [tf.feature_column.numeric_column("x", shape=[2])]
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=tf.train.AdamOptimizer(learning_rate=.001)),
"dnn":
dnn_estimator_fn(
head=head,
feature_columns=feature_columns,
optimizer=tf.train.AdamOptimizer(learning_rate=.001),
hidden_units=[3]),
"dnn2":
dnn_estimator_fn(
head=head,
feature_columns=feature_columns,
optimizer=tf.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 for three iterations.
train_spec = tf.estimator.TrainSpec(input_fn=input_fn, max_steps=300)
eval_spec = tf.estimator.EvalSpec(input_fn=input_fn, steps=1)
# 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)
def test_auto_ensemble_estimator_lifecycle(self):
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
estimator = AutoEnsembleEstimator(
head=head,
candidate_pool=[
linear_estimator_fn(head=head,
feature_columns=feature_columns,
optimizer=optimizer),
dnn_estimator_fn(head=head,
feature_columns=feature_columns,
optimizer=optimizer,
hidden_units=[3]),
],
max_iteration_steps=4,
force_grow=True,
model_dir=self.test_subdirectory,
config=run_config)
# Train.
estimator.train(input_fn=train_input_fn, max_steps=12)
# Evaluate.
eval_results = estimator.evaluate(input_fn=train_input_fn, steps=3)
self.assertIsNotNone(eval_results["loss"])
# 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")
tf.contrib.estimator.export_saved_model_for_mode(
estimator,
export_dir_base=export_dir_base,
input_receiver_fn=serving_input_fn,
mode=tf.estimator.ModeKeys.PREDICT)
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"]))
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
}
head = regression_head.RegressionHead(
loss_reduction=tf_compat.SUM_OVER_BATCH_SIZE)
features = [[1., 0.], [0., 0], [0., 1.], [1., 1.]]
labels = [[1.], [0.], [1.], [0.]]
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])]
candidate_pool = {
"linear":
tf.estimator.LinearEstimator(
head=head,
feature_columns=feature_columns,
optimizer=lambda: tf.keras.optimizers.Adam(lr=.001)),
"dnn":
tf.estimator.DNNEstimator(
head=head,
feature_columns=feature_columns,
optimizer=lambda: tf.keras.optimizers.Adam(lr=.001),
hidden_units=[3]),
"dnn2":
tf.estimator.DNNEstimator(
head=head,
feature_columns=feature_columns,
optimizer=lambda: tf.keras.optimizers.Adam(lr=.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)
elif FLAGS.estimator_type == "autoensemble_trees_multiclass":
n_classes = 3
head = multi_class_head.MultiClassHead(
n_classes=n_classes, loss_reduction=tf_compat.SUM_OVER_BATCH_SIZE)
def tree_loss_fn(labels, logits):
result = bt_losses.per_example_maxent_loss(labels=labels,
logits=logits,
num_classes=n_classes,
weights=None)
return result[0]
tree_head = multi_class_head.MultiClassHead(
loss_fn=tree_loss_fn,
n_classes=n_classes,
loss_reduction=tf_compat.SUM_OVER_BATCH_SIZE)
labels = [[1], [0], [1], [2]]
feature_columns = [tf.feature_column.numeric_column("x", shape=[2])]
candidate_pool = lambda config: { # pylint: disable=g-long-lambda
"linear":
tf.estimator.LinearEstimator(
head=head,
feature_columns=feature_columns,
optimizer=tf.keras.optimizers.Adam(lr=.001),
config=config),
"gbdt":
CoreGradientBoostedDecisionTreeEstimator(
head=tree_head,
learner_config=learner_pb2.LearnerConfig(num_classes=n_classes),
examples_per_layer=8,
num_trees=None,
center_bias=False, # Required for multi-class.
feature_columns=feature_columns,
config=config),
}
estimator = AutoEnsembleEstimator(head=head,
candidate_pool=candidate_pool,
**kwargs)
def input_fn():
input_features = {"x": tf.constant(features, name="x")}
input_labels = tf.constant(labels, name="y")
return input_features, input_labels
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)
