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