tff.learning.programs.extract_and_rewrap_metrics

Extracts a sub-structure and re-wraps it with a new prefix.

tff.learning.programs.extract_and_rewrap_metrics(
    metrics_structure: Mapping[str, Any], *, path: Sequence[str]
) -> Mapping[str, Any]

This is used to normalize outputs from the training and evaluation computations. For example, we are interested in the following substructure from the evaluation and training tasks:

  client_work/eval/current_round_metrics/...
  client_work/eval/total_rounds_metrics/...
  client_work/train/...

And we like to "re-home" these structures under:

  model_metrics/...

while leaving all the other metrics alone. This can be thought of as a "subtree promotion" method.

This is used for grouping metrics so that they appear aligned in http://tensorboard/ in a uniform way. TensorBoard uses the name up to the first / as a tab separator in the UI.

metrics_structure The nested structure of tensor-like metric values.
path A sequence of strings that will be used to traverse the keys of metrics_structure and identify the substructure of interest.

A structure of metrics of metrics_substructure.

ValueError If path is empty.
KeyError If any key in path sequence does not exist in the metrics_structure.