tff.learning.metrics.finalize_then_sample

View source on GitHub

Creates a TFF computation to aggregate metrics via finalize_then_sample.

tff.learning.metrics.finalize_then_sample(
    metric_finalizers: Union[tff.learning.metrics.MetricFinalizersType, tff.learning.metrics.FunctionalMetricFinalizersType],
    local_unfinalized_metrics_type: Optional[tff.types.StructWithPythonType] = None,
    sample_size: int = 100
) -> tff.Computation

The returned federated TFF computation is a polymorphic computation that accepts unfinalized client metrics, and returns finalized, summed metrics placed at the server. Invoking the polymorphic computation will initiate tracing on the argument and will raise a ValueError if the keys (i.e., metric names) in metric_finalizers are not the same as those of the argument the polymorphic method is invoked on.

The returned computation is intended to be invoked on the output of tff.learning.models.VariableModel.report_local_unfinalized_metrics() when placed at CLIENTS. The output is computed by first finalizing each client's metrics locally, and then collecting metrics from at most sample_size clients at the SERVER. If more than sample_size clients participating, then sample_size clients are sampled (by reservoir sampling algorithm); otherwise, all clients' metrics are collected. Sampling is done in a "per-client" manner, i.e., a client, once sampled, will contribute all its metrics to the final result.

The collected metrics samples at SERVER has the same structure (i.e., same keys in a dictionary) as the client's local metrics, except that each leaf node contains a list of scalar metric values, where each value comes from a sampled client, e.g.,

  sampled_metrics_at_server = {
      'metric_a': [a1, a2, ...],
      'metric_b': [b1, b2, ...],
      ...
  }

where "a1" and "b1" are from the same client (similary for "a2" and "b2" etc).

training_process = tff.learning.algorithms.build_weighted_fed_avg(
    model_fn=..., ..., metrics_aggregator=finalize_then_sample,
)
state = training_process.initialize()
for i in range(num_rounds):
  output = training_process.next(state, client_data_at_round_i)
  state = output.state
  sampled_client_metrics = output.metrics['client_work']