def test_ensembler(self, combining_type, search_type, priors, logits,
output_graph):
# Force graph mode
with tf.compat.v1.Graph().as_default():
spec = phoenix_spec_pb2.PhoenixSpec()
spec.ensemble_spec.combining_type = combining_type
spec.ensemble_spec.ensemble_search_type = search_type
ensembler_instance = ensembler.Ensembler(spec)
priors_logits_specs = []
search_logits_specs = []
if priors:
for _ in range(priors):
spec = architecture_utils.LogitsSpec(
logits=tf.zeros([20, 10]))
priors_logits_specs.append(spec)
if logits:
spec = architecture_utils.LogitsSpec(
logits=tf.keras.layers.Dense(10)(tf.zeros([20, 10])))
search_logits_specs.append(spec)
_ = ensembler_instance.bundle_logits(
priors_logits_specs=priors_logits_specs,
search_logits_specs=search_logits_specs,
logits_dimension=10)
nodes = tf.compat.v1.get_default_graph().as_graph_def().node
logging.info([node.name for node in nodes])
self.assertCountEqual([n.name for n in nodes], output_graph)
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 bundle_logits(self, priors_logits_specs, search_logits_specs):
"""Bundles the priors and the search candidate."""
assert search_logits_specs, "Cannot distill with no student model."
assert len(search_logits_specs) == 1, "Search has more than one tower."
if not priors_logits_specs:
return DistillationLogits(
train_logits_specs=search_logits_specs,
eval_logits_spec=search_logits_specs[0],
teacher_logits_spec=None)
with tf.compat.v1.variable_scope("Phoenix/Distiller"):
priors_logits = tf.add_n(
[tf.stop_gradient(spec.logits) for spec in priors_logits_specs])
assert self._distillation_spec.distillation_type, (
"Invalid DistillationType specified.")
if (self._distillation_spec.distillation_type ==
distillation_spec_pb2.DistillationSpec.DistillationType.MSE_LOGITS):
transformed_logits = priors_logits
else:
transformed_logits = tf.nn.softmax(priors_logits /
self._distillation_spec.temperature)
transformed_logits_specs = architecture_utils.LogitsSpec(
logits=transformed_logits)
# Use the logits from the student model (search) to train and evaluate,
# but store the logits from the teacher model (combined priors) to
# calculate the loss.
return DistillationLogits(
train_logits_specs=search_logits_specs,
eval_logits_spec=search_logits_specs[0],
teacher_logits_spec=transformed_logits_specs)
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 create_logits_spec(self,
phoenix_spec,
pre_logits,
dimension,
is_frozen,
lengths=None):
"""Creates the logits for the tower.
Args:
phoenix_spec: The trial's `phoenix_spec_pb2.PhoenixSpec` proto.
pre_logits: `tf.Tensor` of the layer before the logits layer.
dimension: int - the output tensor last axis dimension.
is_frozen: Whether the tower should be frozen.
lengths: A tensor of shape [batch] holding the sequence length for a
sequential problem (rnn).
Returns:
A LogitsSpec containing the main and auxiliary logits and the architecture
of the underlying tower.
"""
logits_weight = 1.0
aux_logits = None
aux_logits_weight = None
if (phoenix_spec.problem_type ==
phoenix_spec_pb2.PhoenixSpec.RNN_ALL_ACTIVATIONS):
logits = tf.compat.v1.layers.conv1d(
inputs=pre_logits, filters=dimension, kernel_size=1)
elif (phoenix_spec.problem_type ==
phoenix_spec_pb2.PhoenixSpec.RNN_LAST_ACTIVATIONS):
if lengths is not None:
logits = utils.last_activations_in_sequence(
tf.compat.v1.layers.conv1d(
inputs=pre_logits, filters=dimension, kernel_size=1), lengths)
else:
logging.warning("Length is missing for rnn_last problem type.")
logits = tf.compat.v1.layers.conv1d(
inputs=pre_logits, filters=dimension, kernel_size=1)
elif phoenix_spec.problem_type in (phoenix_spec_pb2.PhoenixSpec.CNN,
phoenix_spec_pb2.PhoenixSpec.DNN):
logits = tf.keras.layers.Dense(dimension, name="dense")(pre_logits)
else:
raise ValueError("phoenix_spec.problem_type must be either DNN, CNN, "
"RNN_LAST_ACTIVATIONS, or RNN_ALL_ACTIVATIONS.")
logits = tf.identity(logits, name="logits")
if aux_logits is not None:
aux_logits = tf.identity(aux_logits, name="aux_logits")
# TODO(b/172564129): Remove from eval graph.
if is_frozen:
logits = tf.stop_gradient(logits)
if aux_logits is not None:
aux_logits = tf.stop_gradient(aux_logits)
return architecture_utils.LogitsSpec(logits, logits_weight, aux_logits,
aux_logits_weight)
def _create_average_ensemble_logits(self, ensemble_logits,
search_logits_specs):
"""Bundles together averaged logits."""
logits = tf.add_n(ensemble_logits) / len(ensemble_logits)
ensemble_logits_spec = architecture_utils.LogitsSpec(logits=logits)
if (trial_utils.is_nonadaptive_ensemble_search(self._ensemble_spec)
or trial_utils.is_intermixed_ensemble_search(
self._ensemble_spec)):
return EnsembleLogits(train_logits_specs=[],
eval_logits_spec=ensemble_logits_spec)
if trial_utils.is_adaptive_ensemble_search(self._ensemble_spec):
return EnsembleLogits(train_logits_specs=search_logits_specs,
eval_logits_spec=ensemble_logits_spec)
if trial_utils.is_residual_ensemble_search(self._ensemble_spec):
return EnsembleLogits(train_logits_specs=[ensemble_logits_spec],
eval_logits_spec=ensemble_logits_spec)
def _create_weighted_ensemble_logits(self, ensemble_logits,
search_logits_specs,
logits_dimension):
"""Bundles together weighted logits."""
logits = tf.keras.layers.Dense(units=logits_dimension)(tf.concat(
ensemble_logits, axis=-1))
ensemble_logits_spec = architecture_utils.LogitsSpec(logits=logits)
if (trial_utils.is_nonadaptive_ensemble_search(self._ensemble_spec)
or trial_utils.is_intermixed_ensemble_search(
self._ensemble_spec)):
return EnsembleLogits(train_logits_specs=[ensemble_logits_spec],
eval_logits_spec=ensemble_logits_spec)
if trial_utils.is_adaptive_ensemble_search(self._ensemble_spec):
return EnsembleLogits(train_logits_specs=[ensemble_logits_spec] +
search_logits_specs,
eval_logits_spec=ensemble_logits_spec)
if trial_utils.is_residual_ensemble_search(self._ensemble_spec):
return EnsembleLogits(train_logits_specs=[ensemble_logits_spec],
eval_logits_spec=ensemble_logits_spec)
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())