def test_tpu(self):
# Force graph mode
with tf.compat.v1.Graph().as_default():
learning_rate_spec = {
'learning_rate': 0.001,
'gradient_max_norm': 3
}
spec = phoenix_spec_pb2.PhoenixSpec(
problem_type=phoenix_spec_pb2.PhoenixSpec.DNN)
task_manager_instance = task_manager.TaskManager(spec)
logits = tf.keras.layers.Dense(10)(tf.zeros([20, 10]))
logits_spec = architecture_utils.LogitsSpec(logits=logits)
features = {'x': tf.zeros([10, 10])}
loss_fn = loss_fns.make_multi_class_loss_fn()
_ = task_manager_instance.create_model_spec(
features=features,
params=hp.HParams(optimizer='sgd'),
learning_rate_spec=learning_rate_spec,
train_logits_specs=[logits_spec],
eval_logits_spec=logits_spec,
labels=tf.ones([20], dtype=tf.int32),
loss_fn=loss_fn,
mode=tf.estimator.ModeKeys.TRAIN,
lengths=None,
use_tpu=True,
predictions_fn=_default_predictions_fn)
self.assertNotEmpty([
node.name for node in
tf.compat.v1.get_default_graph().as_graph_def().node
if 'CrossReplicaSum' in node.name
])
def test_learning_spec_on_predict(self, learning_rate_spec,
not_containing):
spec = phoenix_spec_pb2.PhoenixSpec(
problem_type=phoenix_spec_pb2.PhoenixSpec.DNN)
task_manager_instance = task_manager.TaskManager(spec)
logits = tf.keras.layers.Dense(10)(tf.zeros([20, 10]))
logits_spec = architecture_utils.LogitsSpec(logits=logits)
features = {'x': tf.zeros([10, 10])}
loss_fn = loss_fns.make_multi_class_loss_fn()
model = task_manager_instance.create_model_spec(
features=features,
params=hp.HParams(optimizer='sgd'),
learning_rate_spec=learning_rate_spec,
train_logits_specs=[logits_spec],
eval_logits_spec=logits_spec,
labels=tf.ones([20], dtype=tf.int32),
loss_fn=loss_fn,
mode=tf.estimator.ModeKeys.PREDICT,
lengths=None,
use_tpu=False,
predictions_fn=_default_predictions_fn)
for phrase in not_containing:
self.assertEmpty([
node.name for node in
tf.compat.v1.get_default_graph().as_graph_def().node
if phrase in node.name
])
self.assertLen(model.predictions, 3)
self.assertIn('probabilities', model.predictions)
self.assertIn('log_probabilities', model.predictions)
self.assertIn('predictions', model.predictions)
self.assertIsNone(model.loss)
def test_architecture(self):
# Force graph mode
with tf.compat.v1.Graph().as_default():
learning_rate_spec = {'learning_rate': 0.001, 'gradient_max_norm': 3}
spec = phoenix_spec_pb2.PhoenixSpec(
problem_type=phoenix_spec_pb2.PhoenixSpec.CNN)
text_format.Merge(
"""
multi_task_spec {
label_name: "label1"
number_of_classes: 10
architecture: "FIXED_OUTPUT_FULLY_CONNECTED_128"
}
multi_task_spec {
label_name: "label2"
number_of_classes: 10
architecture: "FIXED_OUTPUT_FULLY_CONNECTED_256"
architecture: "FIXED_OUTPUT_FULLY_CONNECTED_512"
}
""", spec)
task_manager_instance = task_manager.TaskManager(spec)
logits = tf.keras.layers.Dense(10)(tf.zeros([20, 10]))
logits_spec = architecture_utils.LogitsSpec(logits=logits)
features = {'x': tf.zeros([10, 10])}
loss_fn = loss_fns.make_multi_class_loss_fn()
model = task_manager_instance.create_model_spec(
features=features,
params=hp.HParams(optimizer='sgd'),
learning_rate_spec=learning_rate_spec,
train_logits_specs=[logits_spec],
eval_logits_spec=logits_spec,
labels={
'label1': tf.ones([20], dtype=tf.int32),
'label2': tf.ones([20], dtype=tf.int32)
},
loss_fn=loss_fn,
model_directory=self.get_temp_dir(),
mode=tf.estimator.ModeKeys.TRAIN,
lengths=None,
use_tpu=False,
predictions_fn=_default_predictions_fn)
self.assertNotEmpty([
node.name
for node in tf.compat.v1.get_default_graph().as_graph_def().node
if 'task_label1_tower/1_FIXED_OUTPUT_FULLY_CONNECTED_128' in node.name
])
self.assertNotEmpty([
node.name
for node in tf.compat.v1.get_default_graph().as_graph_def().node
if 'task_label2_tower/1_FIXED_OUTPUT_FULLY_CONNECTED_256' in node.name
])
self.assertNotEmpty([
node.name
for node in tf.compat.v1.get_default_graph().as_graph_def().node
if 'task_label2_tower/2_FIXED_OUTPUT_FULLY_CONNECTED_512' in node.name
])
self.assertLen(model.predictions, 3 * (1 + 2))
self.assertIn('probabilities', model.predictions)
self.assertIn('log_probabilities', model.predictions)
self.assertIn('predictions', model.predictions)
def test_weight_feature(self, is_multitask, weight_is_a_feature):
# Force graph mode
with tf.compat.v1.Graph().as_default():
learning_rate_spec = {
'learning_rate': 0.001,
'gradient_max_norm': 3
}
spec = phoenix_spec_pb2.PhoenixSpec(
problem_type=phoenix_spec_pb2.PhoenixSpec.DNN)
labels = tf.ones([20], dtype=tf.int32)
if is_multitask:
text_format.Merge(
"""
multi_task_spec {
label_name: "label1"
number_of_classes: 10
weight_feature_name: "weight1"
weight_is_a_feature: %s
}
multi_task_spec {
label_name: "label2"
number_of_classes: 10
weight_feature_name: "weight2"
weight_is_a_feature: %s
}
""" % (str(weight_is_a_feature), str(weight_is_a_feature)), spec)
labels = {
'label1': tf.ones([20], dtype=tf.int32),
'label2': tf.ones([20], dtype=tf.int32)
}
weights = {
'weight1': tf.constant([2] * 20),
'weight2': tf.constant([3] * 20)
}
features = {'x': tf.zeros([10, 10])}
if weight_is_a_feature:
features.update(weights)
elif isinstance(labels, dict):
labels.update(weights)
task_manager_instance = task_manager.TaskManager(spec)
logits = tf.keras.layers.Dense(10)(tf.zeros([20, 10]))
logits_spec = architecture_utils.LogitsSpec(logits=logits)
loss_fn = loss_fns.make_multi_class_loss_fn()
_ = task_manager_instance.create_model_spec(
features=features,
params=hp.HParams(optimizer='sgd'),
learning_rate_spec=learning_rate_spec,
train_logits_specs=[logits_spec],
eval_logits_spec=logits_spec,
labels=labels,
loss_fn=loss_fn,
mode=tf.estimator.ModeKeys.TRAIN,
lengths=None,
use_tpu=False,
predictions_fn=_default_predictions_fn)
def test_wrong_dict_weight_feature(self, weight_is_a_feature):
learning_rate_spec = {'learning_rate': 0.001, 'gradient_max_norm': 3}
spec = phoenix_spec_pb2.PhoenixSpec(
problem_type=phoenix_spec_pb2.PhoenixSpec.DNN)
text_format.Merge(
"""
multi_task_spec {
label_name: "label1"
number_of_classes: 10
weight_feature_name: "weight1"
weight_is_a_feature: %s
}
multi_task_spec {
label_name: "label2"
number_of_classes: 10
weight_feature_name: "weight2"
weight_is_a_feature: %s
}
""" % (str(weight_is_a_feature), str(weight_is_a_feature)), spec)
labels = {
'label1': tf.ones([20], dtype=tf.int32),
'label2': tf.ones([20], dtype=tf.int32),
}
# Fix the size of the dict labels to bypass the assertion.
if not weight_is_a_feature:
labels.update({
'not_used': tf.ones([20], dtype=tf.int32),
'not_used2': tf.ones([20], dtype=tf.int32)
})
weights = {
'weight1': tf.constant([2] * 20),
'weight2': tf.constant([3] * 20)
}
features = {'x': tf.zeros([10, 10])}
if not weight_is_a_feature:
features.update(weights)
task_manager_instance = task_manager.TaskManager(spec)
logits = tf.keras.layers.Dense(10)(tf.zeros([20, 10]))
logits_spec = architecture_utils.LogitsSpec(logits=logits)
with self.assertRaises(KeyError):
loss_fn = loss_fns.make_multi_class_loss_fn()
_ = task_manager_instance.create_model_spec(
features=features,
params=hp.HParams(optimizer='sgd'),
learning_rate_spec=learning_rate_spec,
train_logits_specs=[logits_spec],
eval_logits_spec=logits_spec,
labels=labels,
loss_fn=loss_fn,
model_directory=self.get_temp_dir(),
mode=tf.estimator.ModeKeys.TRAIN,
lengths=None,
use_tpu=False,
predictions_fn=_default_predictions_fn)
def test_multitask(self, learning_rate_spec, contains_node,
not_containing):
# Force graph mode
with tf.compat.v1.Graph().as_default():
spec = phoenix_spec_pb2.PhoenixSpec(
problem_type=phoenix_spec_pb2.PhoenixSpec.DNN)
text_format.Merge(
"""
multi_task_spec {
label_name: "label1"
number_of_classes: 10
}
multi_task_spec {
label_name: "label2"
number_of_classes: 10
}
""", spec)
task_manager_instance = task_manager.TaskManager(spec)
logits = tf.keras.layers.Dense(10)(tf.zeros([20, 10]))
logits_spec = architecture_utils.LogitsSpec(logits=logits)
features = {'x': tf.zeros([10, 10])}
loss_fn = loss_fns.make_multi_class_loss_fn()
model = task_manager_instance.create_model_spec(
features=features,
params=hp.HParams(optimizer='sgd'),
learning_rate_spec=learning_rate_spec,
train_logits_specs=[logits_spec],
eval_logits_spec=logits_spec,
labels={
'label1': tf.ones([20], dtype=tf.int32),
'label2': tf.ones([20], dtype=tf.int32)
},
loss_fn=loss_fn,
mode=tf.estimator.ModeKeys.TRAIN,
lengths=None,
use_tpu=False,
predictions_fn=_default_predictions_fn)
self.assertNotEmpty([
node.name for node in
tf.compat.v1.get_default_graph().as_graph_def().node
if contains_node in node.name
])
for phrase in not_containing:
self.assertEmpty([
node.name for node in
tf.compat.v1.get_default_graph().as_graph_def().node
if phrase in node.name
])
self.assertLen(model.predictions, 3 * (1 + 2))
self.assertContainsSubset([
'probabilities',
'probabilities/label1',
'probabilities/label2',
'log_probabilities',
'log_probabilities/label1',
'log_probabilities/label2',
'predictions',
'predictions/label1',
'predictions/label2',
], model.predictions.keys())
def __init__(self,
phoenix_spec,
input_layer_fn,
study_owner,
study_name,
head=None,
logits_dimension=None,
label_vocabulary=None,
loss_fn=None,
metric_fn=None,
predictions_fn=None,
metadata=None):
"""Constructs a Phoenix instance.
Args:
phoenix_spec: A `PhoenixSpec` proto with the spec for the run.
input_layer_fn: A function that converts feature Tensors to input layer.
See learning.autolx.model_search.data.Provider.get_input_layer_fn
for details.
study_owner: A string holding the ldap of the study owner. We use tuner
platforms to conduct the various architectures training. This field
specifies the study owner.
study_name: A string holding the study name.
head: A head to use with Phoenix for creating the loss and eval metrics.
If no head is given, Phoenix falls back to using the loss_fn and
metric_fn. N.B.: Phoenix creates its own EstimatorSpec so everything
besides the loss and eval metrics returned by head will be ignored.
logits_dimension: An int holding the dimension of the output. Must be
provided if head is None. Will be ignored if head is not None.
label_vocabulary: List or tuple with labels vocabulary. Needed only if the
labels are of type string. This list is used by the loss function if
loss_fn is not provided. It is also used in the metric function to
create the accuracy metric ops. Use only with multiclass classification
problems.
loss_fn: A function to compute the loss. Ignored if `head` is not None.
Must accept as inputs a `labels` Tensor, a `logits` Tensor, and
optionally a `weights` Tensor. `weights` must either be rank 0 or have
the same rank as labels. If None, Phoenix defaults to using softmax
cross-entropy.
metric_fn: Metrics for Tensorboard. Ignored if `head` is not None.
metric_fn takes `label` and `predictions` as input, and outputs a
dictionary of (tensor, update_op) tuples. `label` is a Tensor (in the
single task case) or a dict of Tensors (in the case of multi-task, where
the key of the dicts correspond to the task names). `predictions` is a
dict of Tensors. In the single task case, it consists of `predictions`,
`probabilities`, and `log_probabilities`. In the multi-task case, it
consists of the same keys as that of the single task case, but also
those corresponding to each task (e.g., predictions/task_name_1). See
`metric_fns` for more detail. If `metric_fn` is None, it will include a
metric for the number of parameters, accuracy (if logit_dimensions >=
2), and AUC metrics (if logit_dimensions == 2).
predictions_fn: A function to convert eval logits to the
`predictions` dictionary passed to metric_fn. If `None`, defaults to
computing 'predictions', 'probabilities', and 'log_probabilities'.
metadata: An object that implements metadata api in
learning.adanets.phoenix.metadata.Metadata
"""
# Check Phoenix preconditions and fail early if any of them are broken.
if phoenix_spec.multi_task_spec:
# TODO(b/172564129): Add support for head and custom loss_fns in
# multi-task.
assert not head, "head is not supported for multi-task."
if head:
msg = "Do not specify {} when using head as head already contains it."
assert not logits_dimension, msg.format("logits_dimension")
assert not label_vocabulary, msg.format("label_vocabulary")
assert not loss_fn, msg.format("loss_fn")
assert not metric_fn, msg.format("metric_fn")
# Check ensemble search / distillation preconditions.
ensemble_spec = phoenix_spec.ensemble_spec
distillation_spec = phoenix_spec.distillation_spec
if trial_utils.has_distillation(
distillation_spec) and trial_utils.has_ensemble_search(
ensemble_spec
) and not trial_utils.is_intermixed_ensemble_search(ensemble_spec):
ensemble_search_spec = (
ensemble_spec.nonadaptive_search
if trial_utils.is_nonadaptive_ensemble_search(ensemble_spec) else
ensemble_spec.adaptive_search)
if (distillation_spec.minimal_pool_size ==
ensemble_search_spec.minimal_pool_size):
logging.warning("minimal_pool_size is the same for ensemble spec and "
"distillation spec, so distillation will be ignored.")
self._phoenix_spec = phoenix_spec
self._input_layer_fn = input_layer_fn
self._ensembler = ensembler.Ensembler(phoenix_spec)
self._distiller = distillation.Distiller(phoenix_spec.distillation_spec)
self._study_owner = study_owner
self._study_name = study_name
self._head = head
self._logits_dimension = (
self._head.logits_dimension if head else logits_dimension)
self._label_vocabulary = label_vocabulary
if self._label_vocabulary:
assert self._logits_dimension == len(self._label_vocabulary)
self._loss_fn = loss_fn or loss_fns.make_multi_class_loss_fn(
label_vocabulary=label_vocabulary)
self._user_specified_metric_fn = metric_fn
self._predictions_fn = (predictions_fn or _default_predictions_fn)
if metadata is None:
self._metadata = ml_metadata_db.MLMetaData(phoenix_spec, study_name,
study_owner)
else:
self._metadata = metadata
self._task_manager = task_manager.TaskManager(phoenix_spec)
self._controller = controller.InProcessController(
phoenix_spec=phoenix_spec, metadata=self._metadata)
