def _prefix(tensors, group_name, flat_key, default_key):
"""Prefixes tensors by group_name and either flat_key or default_key.
If tensors is a dict each tensor is rekeyed as group_name/flat_key/key. If
tensors is a single Tensor, it is keyed by group_name/default_key.
Args:
tensors: A Tensor or dictionary of Tensors.
group_name: The group name to use in the prefix.
flat_key: The key to use in the prefix if tensors is a dictionary.
default_key: The default key to use if tensors is a single Tensor.
Returns:
A dictionary of tensors prefixed by group_name and key. If tensors is a
single Tensor, the returned dictionary will only have one element.
"""
prefix = default_key
if isinstance(tensors, dict):
prefix = flat_key
tensors_ = {}
for key in six.iterkeys(tensors):
# multi_head uses tuples of strings as the key.
if isinstance(key, tuple):
tensors_["|".join(key)] = tensors[key]
else:
tensors_[key] = tensors[key]
tensors = tensors_
tensors = {prefix: tensors}
tensors = dict_utils.flatten_dict(tensors)
tensors = dict_utils.flatten_dict({group_name: tensors})
return tensors
def test_flatten_dict(self):
to_flatten = {
"hello": {
"world": 1,
"sailor": 2,
},
"ada": {
"net": 3,
"boost": 4,
},
"nodict": 5,
}
actual = dict_utils.flatten_dict(to_flatten, delimiter="-")
expected = {
"hello-world": 1,
"hello-sailor": 2,
"ada-net": 3,
"ada-boost": 4,
"nodict": 5,
}
self.assertDictEqual(actual, expected)
def _create_best_eval_metrics_tuple(self, candidates, subnetwork_specs,
best_candidate_index, mode, params):
"""Returns (metric_fn, tensors) which computes the best ensemble's metrics.
Specifically, when metric_fn(tensors) is called, it separates the metric ops
by metric name. All candidates are not required to have the same metrics.
When they all share a given metric, an additional metric is added which
represents that of the best candidate.
Args:
candidates: List of `_Candidate` instances to choose from.
subnetwork_specs: List of `_SubnetworkSpec` instances for this iteration.
best_candidate_index: `Tensor` index of the best candidate in the list.
mode: Defines whether this is training, evaluation or inference. Eval
metrics are only defined during evaluation. See `ModeKeys`.
params: The params passed to model_fn.
Returns:
Dict of metric results keyed by name. The values of the dict are the
results of calling a metric function.
"""
if mode != tf.estimator.ModeKeys.EVAL:
return None
metric_fns, tensors = self._collate_metric_fns_and_tensors(
candidates, subnetwork_specs)
# All tensors outfed from the TPU must be batch-major.
batch_size = params.get("batch_size", 1) if params else 1
tensors["best_candidate_index"] = tf.tile([best_candidate_index],
[batch_size])
tensors = dict_utils.flatten_dict(tensors)
def _best_eval_metrics_fn(**kwargs):
"""Returns the best eval metrics."""
with tf.variable_scope("best_eval_metrics"):
subnetwork_metric_fns = {
k: metric_fns[k]
for k in metric_fns if k.startswith("subnetwork_")
}
subnetwork_tensors = dict_utils.unflatten_dict(
kwargs, subnetwork_metric_fns.keys())
subnetwork_metric_ops = self._group_metric_ops(
subnetwork_metric_fns, subnetwork_tensors)
ensemble_metric_fns = {
k: metric_fns[k]
for k in metric_fns if k.startswith("ensemble_")
}
ensemble_tensors = dict_utils.unflatten_dict(
kwargs, ensemble_metric_fns.keys())
grouped_metrics = self._group_metric_ops(
ensemble_metric_fns, ensemble_tensors)
eval_metric_ops = {}
for metric_name in sorted(grouped_metrics):
metric_ops = grouped_metrics[metric_name]
if len(metric_ops) != len(candidates):
continue
if metric_name == "loss":
continue
best_candidate_index = kwargs["best_candidate_index"]
values, ops = list(six.moves.zip(*metric_ops))
idx, idx_update_op = tf.metrics.mean(best_candidate_index)
best_value = tf.stack(values)[tf.cast(idx, tf.int32)]
# All tensors in this function have been outfed from the TPU, so we
# must update them manually, otherwise the TPU will hang indefinetly
# for the value of idx to update.
ops = list(ops)
ops.append(idx_update_op)
# Bundle subnetwork eval metric ops and ensemble "loss"" ops (which
# is a restricted Estimator keyword) into other metric ops so that
# they are computed.
ensemble_loss_ops = grouped_metrics.get("loss", tf.no_op())
all_ops = tf.group(ops, ensemble_loss_ops,
subnetwork_metric_ops)
eval_metric_ops[metric_name] = (best_value, all_ops)
# tf.estimator.Estimator does not allow a "loss" key to be present in
# its eval_metrics.
assert "loss" not in eval_metric_ops
return eval_metric_ops
return _best_eval_metrics_fn, tensors
def _create_best_eval_metrics_tuple(self, candidates, best_candidate_index,
mode, params):
"""Returns (metric_fn, tensors) which computes best candidate metrics.
Specifically, when metric_fn(tensors) is called, it separates the metric ops
by metric name. All candidates are not required to have the same metrics.
When they all share a given metric, an additional metric is added which
represents that of the best candidate.
Args:
candidates: List of `_Candidate` instances to choose from.
best_candidate_index: `Tensor` index of the best candidate in the list.
mode: Defines whether this is training, evaluation or inference. Eval
metrics are only defined during evaluation. See `ModeKeys`.
params: The params passed to model_fn.
Returns:
Dict of metric results keyed by name. The values of the dict are the
results of calling a metric function.
"""
if mode != tf.estimator.ModeKeys.EVAL:
return None
metric_fns, tensors = self._collate_metric_fns_and_tensors(candidates)
# All tensors outfed from the TPU must be batch-major.
batch_size = params.get("batch_size", 1) if params else 1
tensors["best_candidate_index"] = tf.tile([best_candidate_index],
[batch_size])
tensors = dict_utils.flatten_dict(tensors)
def _best_eval_metrics_fn(**kwargs):
"""Returns the best eval metrics."""
with tf.variable_scope("best_eval_metrics"):
tensors = dict_utils.unflatten_dict(kwargs, metric_fns.keys())
grouped_metrics = self._group_metric_ops(metric_fns, tensors)
eval_metric_ops = {}
for metric_name in sorted(grouped_metrics):
metric_ops = grouped_metrics[metric_name]
if len(metric_ops) != len(candidates):
continue
best_candidate_index = tensors["best_candidate_index"]
values, ops = list(six.moves.zip(*metric_ops))
idx, idx_update_op = tf.metrics.mean(best_candidate_index)
best_value = tf.stack(values)[tf.cast(idx, tf.int32)]
# All tensors in this function have been outfed from the TPU, so we
# must update them manually, otherwise the TPU will hang indefinetly
# for the value of idx to wait.
ops = list(ops)
ops.append(idx_update_op)
best_op = tf.group(ops)
best_candidate_metric = (best_value, best_op)
eval_metric_ops[metric_name] = best_candidate_metric
# Include any evaluation metric shared among all the candidates in
# the top level metrics in TensorBoard. These "root" metrics track
# AdaNet's overall performance, making it easier to compare with other
# estimators that report the same metrics.
suffix = "/adanet/adanet_weighted_ensemble"
if not metric_name.endswith(suffix):
continue
root_metric_name = metric_name[:-len(suffix)]
if root_metric_name == "loss":
continue
eval_metric_ops[root_metric_name] = best_candidate_metric
return eval_metric_ops
return _best_eval_metrics_fn, tensors