def build_federated_averaging_process(
model_fn: Callable[[], model_lib.Model],
client_optimizer_fn: Callable[[], tf.keras.optimizers.Optimizer],
server_optimizer_fn: Callable[
[], tf.keras.optimizers.Optimizer] = DEFAULT_SERVER_OPTIMIZER_FN,
client_weight_fn: Callable[[Any], tf.Tensor] = None,
stateful_delta_aggregate_fn=None,
stateful_model_broadcast_fn=None) -> tff.utils.IterativeProcess:
"""Builds the TFF computations for optimization using federated averaging.
Args:
model_fn: A no-arg function that returns a `tff.learning.Model`.
client_optimizer_fn: A no-arg callable that returns a `tf.keras.Optimizer`.
server_optimizer_fn: A no-arg callable that returns a `tf.keras.Optimizer`.
The `apply_gradients` method of this optimizer is used to apply client
updates to the server model. The default creates a
`tf.keras.optimizers.SGD` with a learning rate of 1.0, which simply adds
the average client delta to the server's model.
client_weight_fn: Optional function that takes the output of
`model.report_local_outputs` and returns a tensor that provides the weight
in the federated average of model deltas. If not provided, the default is
the total number of examples processed on device.
stateful_delta_aggregate_fn: A `tff.utils.StatefulAggregateFn` where the
`next_fn` performs a federated aggregation and upates state. That is, it
has TFF type `(state@SERVER, value@CLIENTS, weights@CLIENTS) ->
(state@SERVER, aggregate@SERVER)`, where the `value` type is
`tff.learning.framework.ModelWeights.trainable` corresponding to the
object returned by `model_fn`. By default performs arithmetic mean
aggregation, weighted by `client_weight_fn`.
stateful_model_broadcast_fn: A `tff.utils.StatefulBroadcastFn` where the
`next_fn` performs a federated broadcast and upates state. That is, it has
TFF type `(state@SERVER, value@SERVER) -> (state@SERVER, value@CLIENTS)`,
where the `value` type is `tff.learning.framework.ModelWeights`
corresponding to the object returned by `model_fn`. By default performs
identity broadcast.
Returns:
A `tff.utils.IterativeProcess`.
"""
def client_fed_avg(model_fn):
return _ClientFedAvg(model_fn(), client_optimizer_fn(), client_weight_fn)
if stateful_delta_aggregate_fn is None:
stateful_delta_aggregate_fn = optimizer_utils.build_stateless_mean()
else:
py_typecheck.check_type(stateful_delta_aggregate_fn,
tff.utils.StatefulAggregateFn)
if stateful_model_broadcast_fn is None:
stateful_model_broadcast_fn = optimizer_utils.build_stateless_broadcaster()
else:
py_typecheck.check_type(stateful_model_broadcast_fn,
tff.utils.StatefulBroadcastFn)
return optimizer_utils.build_model_delta_optimizer_process(
model_fn, client_fed_avg, server_optimizer_fn,
stateful_delta_aggregate_fn, stateful_model_broadcast_fn)
def test_fails_stateful_broadcast_and_process(self):
model_weights_type = model_utils.weights_type_from_model(
model_examples.LinearRegression)
with self.assertRaises(optimizer_utils.DisjointArgumentError):
federated_averaging.build_federated_averaging_process(
model_fn=model_examples.LinearRegression,
client_optimizer_fn=tf.keras.optimizers.SGD,
stateful_model_broadcast_fn=tff.utils.StatefulBroadcastFn(
initialize_fn=lambda: (),
next_fn=lambda state, weights: # pylint: disable=g-long-lambda
(state, tff.federated_broadcast(weights))),
broadcast_process=optimizer_utils.build_stateless_broadcaster(
model_weights_type=model_weights_type))
def test_fails_stateful_broadcast_and_process(self):
model_weights_type = model_utils.weights_type_from_model(
model_examples.LinearRegression)
with self.assertRaises(optimizer_utils.DisjointArgumentError):
optimizer_utils.build_model_delta_optimizer_process(
model_fn=model_examples.LinearRegression,
model_to_client_delta_fn=DummyClientDeltaFn,
server_optimizer_fn=tf.keras.optimizers.SGD,
stateful_model_broadcast_fn=computation_utils.StatefulBroadcastFn(
initialize_fn=lambda: (),
next_fn=lambda state, weights: # pylint: disable=g-long-lambda
(state, intrinsics.federated_broadcast(weights))),
broadcast_process=optimizer_utils.build_stateless_broadcaster(
model_weights_type=model_weights_type))
def test_fails_stateful_broadcast_and_process(self):
with tf.Graph().as_default():
model_weights_type = tff.framework.type_from_tensors(
model_utils.ModelWeights.from_model(
model_examples.LinearRegression()))
with self.assertRaises(optimizer_utils.DisjointArgumentError):
optimizer_utils.build_model_delta_optimizer_process(
model_fn=model_examples.LinearRegression,
model_to_client_delta_fn=DummyClientDeltaFn,
server_optimizer_fn=tf.keras.optimizers.SGD,
stateful_model_broadcast_fn=tff.utils.StatefulBroadcastFn(
initialize_fn=lambda: (),
next_fn=lambda state, weights: # pylint: disable=g-long-lambda
(state, tff.federated_broadcast(weights))),
broadcast_process=optimizer_utils.build_stateless_broadcaster(
model_weights_type=model_weights_type))
def build_federated_process_for_test(model_fn, num_passes=5, tolerance=1e-6):
"""Build a test FedAvg process with a dummy client computation.
Analogue of `build_federated_averaging_process`, but with client_fed_avg
replaced by the dummy mean computation defined above.
Args:
model_fn: callable that returns a `tff.learning.Model`.
num_passes: integer number of communication rounds in the smoothed
Weiszfeld algorithm (min. 1).
tolerance: float smoothing parameter of smoothed Weiszfeld algorithm.
Default 1e-6.
Returns:
A `tff.utils.IterativeProcess`.
"""
server_optimizer_fn = lambda: tf.keras.optimizers.SGD(learning_rate=1.0)
def client_fed_avg(model_fn):
return DummyClientComputation(model_fn(), client_weight_fn=None)
# Build robust aggregation function
with tf.Graph().as_default():
# workaround since keras automatically appends "_n" to the nth call of
# `model_fn`
model_type = tff.framework.type_from_tensors(
model_fn().weights.trainable)
stateful_delta_aggregate_fn = rfa.build_stateless_robust_aggregation(
model_type,
num_communication_passes=num_passes,
tolerance=tolerance)
stateful_model_broadcast_fn = optimizer_utils.build_stateless_broadcaster(
)
return optimizer_utils.build_model_delta_optimizer_process(
model_fn, client_fed_avg, server_optimizer_fn,
stateful_delta_aggregate_fn, stateful_model_broadcast_fn)
def build_federated_averaging_process(
model_fn: Callable[[], model_lib.Model],
client_optimizer_fn: Callable[[], tf.keras.optimizers.Optimizer],
server_optimizer_fn: Callable[
[], tf.keras.optimizers.Optimizer] = DEFAULT_SERVER_OPTIMIZER_FN,
client_weight_fn: Callable[[Any], tf.Tensor] = None,
stateful_delta_aggregate_fn=None,
stateful_model_broadcast_fn=None) -> tff.utils.IterativeProcess:
"""Builds an iterative process that performs federated averaging.
This function creates a `tff.utils.IterativeProcess` that performs federated
averaging on client models. The iterative process has the following methods:
* `initialize`: A `tff.Computation` with the functional type signature
`( -> S@SERVER)`, where `S` is a`tff.learning.framework.ServerState`
representing the initial state of the server.
* `next`: A `tff.Computation` with the functional type signature
`(<S@SERVER, {B*}@CLIENTS> -> <S@SERVER, T@SERVER>)` where `S` is a
`tff.learning.framework.ServerState` whose type matches that of the output
of `initialize`, and `{B*}@CLIENTS` represents the client datasets, where
`B` is the type of a single batch. This computation returns a
`tff.learning.framework.ServerState` representing the updated server state
and training metrics that are the result of
`tff.learning.Model.federated_output_computation` during client training.
Each time the `next` method is called, the server model is broadcast to each
client using a broadcast function. For each client, one epoch of local
training is performed via the `tf.keras.optimizers.Optimizer.apply_gradients`
method of the client optimizer. Each client computes the difference between
the client model after training and the initial broadcast model. These model
deltas are then aggregated at the server using some aggregation function. The
aggregate model delta is applied at the server by using the
`tf.keras.optimizers.Optimizer.apply_gradients` method of the server
optimizer.
Note: the default server optimizer function is `tf.keras.optimizers.SGD`
with a learning rate of 1.0, which corresponds to adding the model delta to
the current server model. This recovers the original FedAvg algorithm in
[McMahan et al., 2017](https://arxiv.org/abs/1602.05629). More
sophisticated federated averaging procedures may use different learning rates
or server optimizers.
Args:
model_fn: A no-arg function that returns a `tff.learning.Model`.
client_optimizer_fn: A no-arg callable that returns a `tf.keras.Optimizer`.
server_optimizer_fn: A no-arg callable that returns a `tf.keras.Optimizer`.
By default, this uses `tf.keras.optimizers.SGD` with a learning rate of
1.0.
client_weight_fn: Optional function that takes the output of
`model.report_local_outputs` and returns a tensor providing the weight in
the federated average of model deltas. If not provided, the default is the
total number of examples processed on device.
stateful_delta_aggregate_fn: A `tff.utils.StatefulAggregateFn` where the
`next_fn` performs a federated aggregation and upates state. It must have
TFF type `(<state@SERVER, value@CLIENTS, weights@CLIENTS> ->
<state@SERVER, aggregate@SERVER>)`, where the `value` type is
`tff.learning.framework.ModelWeights.trainable` corresponding to the
object returned by `model_fn`. By default performs arithmetic mean
aggregation, weighted by `client_weight_fn`.
stateful_model_broadcast_fn: A `tff.utils.StatefulBroadcastFn` where the
`next_fn` performs a federated broadcast and upates state. It must have
TFF type `(<state@SERVER, value@SERVER> -> <state@SERVER,
value@CLIENTS>)`, where the `value` type is
`tff.learning.framework.ModelWeights` corresponding to the object returned
by `model_fn`. The default is the identity broadcast.
Returns:
A `tff.utils.IterativeProcess`.
"""
def client_fed_avg(model_fn):
return ClientFedAvg(model_fn(), client_optimizer_fn(),
client_weight_fn)
if stateful_delta_aggregate_fn is None:
stateful_delta_aggregate_fn = optimizer_utils.build_stateless_mean()
else:
py_typecheck.check_type(stateful_delta_aggregate_fn,
tff.utils.StatefulAggregateFn)
if stateful_model_broadcast_fn is None:
stateful_model_broadcast_fn = optimizer_utils.build_stateless_broadcaster(
)
else:
py_typecheck.check_type(stateful_model_broadcast_fn,
tff.utils.StatefulBroadcastFn)
return optimizer_utils.build_model_delta_optimizer_process(
model_fn, client_fed_avg, server_optimizer_fn,
stateful_delta_aggregate_fn, stateful_model_broadcast_fn)
def build_federated_sgd_process(
model_fn,
server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.1),
client_weight_fn=None,
stateful_delta_aggregate_fn=None,
stateful_model_broadcast_fn=None):
"""Builds the TFF computations for optimization using federated SGD.
This function creates a `tff.templates.IterativeProcess` that performs
federated averaging on client models. The iterative process has the following
methods:
* `initialize`: A `tff.Computation` with the functional type signature
`( -> S@SERVER)`, where `S` is a`tff.learning.framework.ServerState`
representing the initial state of the server.
* `next`: A `tff.Computation` with the functional type signature
`(<S@SERVER, {B*}@CLIENTS> -> <S@SERVER, T@SERVER>)` where `S` is a
`tff.learning.framework.ServerState` whose type matches that of the output
of `initialize`, and `{B*}@CLIENTS` represents the client datasets, where
`B` is the type of a single batch. This computation returns a
`tff.learning.framework.ServerState` representing the updated server state
and training metrics that are the result of
`tff.learning.Model.federated_output_computation` during client training.
Each time the `next` method is called, the server model is broadcast to each
client using a broadcast function. Each client sums the gradients at each
batch in the client's local dataset. These gradient sums are then aggregated
at the server using an aggregation function. The aggregate gradients are
applied at the server by using the
`tf.keras.optimizers.Optimizer.apply_gradients` method of the server
optimizer.
Note: the default server optimizer function is `tf.keras.optimizers.SGD`
with a learning rate of 1.0, which corresponds to adding the aggregate of the
gradients to the current server model. This recovers the original FedSGD
algorithm in [McMahan et al., 2017](https://arxiv.org/abs/1602.05629). More
sophisticated federated SGD procedures may use different learning rates
or server optimizers.
Args:
model_fn: A no-arg function that returns a `tff.learning.Model`. This method
must *not* capture TensorFlow tensors or variables and use them. The model
must be constructed entirely from scratch on each invocation, returning
the same pre-constructed model each call will result in an error.
server_optimizer_fn: A no-arg function that returns a `tf.Optimizer`. The
`apply_gradients` method of this optimizer is used to apply client updates
to the server model.
client_weight_fn: Optional function that takes the output of
`model.report_local_outputs` and returns a tensor that provides the weight
in the federated average of the aggregated gradients. If not provided, the
default is the total number of examples processed on device.
stateful_delta_aggregate_fn: A `tff.utils.StatefulAggregateFn` where the
`next_fn` performs a federated aggregation and upates state. It must have
TFF type `(<state@SERVER, value@CLIENTS, weights@CLIENTS> ->
<state@SERVER, aggregate@SERVER>)`, where the `value` type is
`tff.learning.framework.ModelWeights.trainable` corresponding to the
object returned by `model_fn`. By default performs arithmetic mean
aggregation, weighted by `client_weight_fn`.
stateful_model_broadcast_fn: A `tff.utils.StatefulBroadcastFn` where the
`next_fn` performs a federated broadcast and upates state. It must have
TFF type `(<state@SERVER, value@SERVER> -> <state@SERVER,
value@CLIENTS>)`, where the `value` type is
`tff.learning.framework.ModelWeights` corresponding to the object returned
by `model_fn`. The default is the identity broadcast.
Returns:
A `tff.templates.IterativeProcess`.
"""
def client_sgd_avg(model_fn):
return ClientSgd(model_fn(), client_weight_fn)
if stateful_delta_aggregate_fn is None:
stateful_delta_aggregate_fn = optimizer_utils.build_stateless_mean()
else:
py_typecheck.check_type(stateful_delta_aggregate_fn,
tff.utils.StatefulAggregateFn)
if stateful_model_broadcast_fn is None:
stateful_model_broadcast_fn = optimizer_utils.build_stateless_broadcaster(
)
else:
py_typecheck.check_type(stateful_model_broadcast_fn,
tff.utils.StatefulBroadcastFn)
return optimizer_utils.build_model_delta_optimizer_process(
model_fn, client_sgd_avg, server_optimizer_fn,
stateful_delta_aggregate_fn, stateful_model_broadcast_fn)
def build_federated_evaluation(
model_fn: training_process.ModelFn,
*, # Callers pass below args by name.
loss_fn: training_process.LossFn,
metrics_fn: Optional[training_process.MetricsFn] = None,
reconstruction_optimizer_fn: training_process.OptimizerFn = functools.
partial(tf.keras.optimizers.SGD, 0.1),
dataset_split_fn: Optional[reconstruction_utils.DatasetSplitFn] = None,
broadcast_process: Optional[measured_process_lib.MeasuredProcess] = None,
) -> computation_base.Computation:
"""Builds a `tff.Computation` for evaluating a reconstruction `Model`.
The returned computation proceeds in two stages: (1) reconstruction and (2)
evaluation. During the reconstruction stage, local variables are reconstructed
by freezing global variables and training using `reconstruction_optimizer_fn`.
During the evaluation stage, the reconstructed local variables and global
variables are evaluated using the provided `loss_fn` and `metrics_fn`.
Usage of returned computation:
eval_comp = build_federated_evaluation(...)
metrics = eval_comp(tff.learning.reconstruction.get_global_variables(model),
federated_data)
Args:
model_fn: A no-arg function that returns a
`tff.learning.reconstruction.Model`. This method must *not* capture
Tensorflow tensors or variables and use them. Must be constructed entirely
from scratch on each invocation, returning the same pre-constructed model
each call will result in an error.
loss_fn: A no-arg function returning a `tf.keras.losses.Loss` to use to
reconstruct and evaluate the model. The loss will be applied to the
model's outputs during the evaluation stage. The final loss metric is the
example-weighted mean loss across batches (and across clients).
metrics_fn: A no-arg function returning a list of `tf.keras.metrics.Metric`s
to evaluate the model. The metrics will be applied to the model's outputs
during the evaluation stage. Final metric values are the example-weighted
mean of metric values across batches (and across clients). If None, no
metrics are applied.
reconstruction_optimizer_fn: A no-arg function that returns a
`tf.keras.optimizers.Optimizer` used to reconstruct the local variables
with the global ones frozen.
dataset_split_fn: A `tff.learning.reconstruction.DatasetSplitFn` taking in a
single TF dataset and producing two TF datasets. The first is iterated
over during reconstruction, and the second is iterated over during
evaluation. This can be used to preprocess datasets to e.g. iterate over
them for multiple epochs or use disjoint data for reconstruction and
evaluation. If None, split client data in half for each user, using one
half for reconstruction and the other for evaluation. See
`tff.learning.reconstruction.build_dataset_split_fn` for options.
broadcast_process: A `tff.templates.MeasuredProcess` that broadcasts the
model weights on the server to the clients. It must support the signature
`(input_values@SERVER -> output_values@CLIENT)` and have empty state. If
set to default None, the server model is broadcast to the clients using
the default `tff.federated_broadcast`.
Raises:
TypeError: if `broadcast_process` does not have the expected signature or
has non-empty state.
Returns:
A `tff.Computation` that accepts global model parameters and federated data
and returns example-weighted evaluation loss and metrics.
"""
# Construct the model first just to obtain the metadata and define all the
# types needed to define the computations that follow.
with tf.Graph().as_default():
model = model_fn()
global_weights = reconstruction_utils.get_global_variables(model)
model_weights_type = type_conversions.type_from_tensors(global_weights)
batch_type = computation_types.to_type(model.input_spec)
metrics = [keras_utils.MeanLossMetric(loss_fn())]
if metrics_fn is not None:
metrics.extend(metrics_fn())
federated_output_computation = (
keras_utils.federated_output_computation_from_metrics(metrics))
# Remove unneeded variables to avoid polluting namespace.
del model
del global_weights
del metrics
if dataset_split_fn is None:
dataset_split_fn = reconstruction_utils.build_dataset_split_fn(
split_dataset=True)
if broadcast_process is None:
broadcast_process = optimizer_utils.build_stateless_broadcaster(
model_weights_type=model_weights_type)
if not optimizer_utils.is_valid_broadcast_process(broadcast_process):
raise TypeError(
'broadcast_process type signature does not conform to expected '
'signature (<state@S, input@S> -> <state@S, result@C, measurements@S>).'
' Got: {t}'.format(t=broadcast_process.next.type_signature))
if iterative_process.is_stateful(broadcast_process):
raise TypeError(
f'Eval broadcast_process must be stateless (have an empty '
'state), has state '
f'{broadcast_process.initialize.type_signature.result!r}')
@tensorflow_computation.tf_computation(
model_weights_type, computation_types.SequenceType(batch_type))
def client_computation(incoming_model_weights: computation_types.Type,
client_dataset: computation_types.SequenceType):
"""Reconstructs and evaluates with `incoming_model_weights`."""
client_model = model_fn()
client_global_weights = reconstruction_utils.get_global_variables(
client_model)
client_local_weights = reconstruction_utils.get_local_variables(
client_model)
metrics = [keras_utils.MeanLossMetric(loss_fn())]
if metrics_fn is not None:
metrics.extend(metrics_fn())
client_loss = loss_fn()
reconstruction_optimizer = reconstruction_optimizer_fn()
@tf.function
def reconstruction_reduce_fn(num_examples_sum, batch):
"""Runs reconstruction training on local client batch."""
with tf.GradientTape() as tape:
output = client_model.forward_pass(batch, training=True)
batch_loss = client_loss(y_true=output.labels,
y_pred=output.predictions)
gradients = tape.gradient(batch_loss,
client_local_weights.trainable)
reconstruction_optimizer.apply_gradients(
zip(gradients, client_local_weights.trainable))
return num_examples_sum + output.num_examples
@tf.function
def evaluation_reduce_fn(num_examples_sum, batch):
"""Runs evaluation on client batch without training."""
output = client_model.forward_pass(batch, training=False)
# Update each metric.
for metric in metrics:
metric.update_state(y_true=output.labels,
y_pred=output.predictions)
return num_examples_sum + output.num_examples
@tf.function
def tf_client_computation(incoming_model_weights, client_dataset):
"""Reconstructs and evaluates with `incoming_model_weights`."""
recon_dataset, eval_dataset = dataset_split_fn(client_dataset)
# Assign incoming global weights to `client_model` before reconstruction.
tf.nest.map_structure(lambda v, t: v.assign(t),
client_global_weights,
incoming_model_weights)
recon_dataset.reduce(tf.constant(0), reconstruction_reduce_fn)
eval_dataset.reduce(tf.constant(0), evaluation_reduce_fn)
eval_local_outputs = keras_utils.read_metric_variables(metrics)
return eval_local_outputs
return tf_client_computation(incoming_model_weights, client_dataset)
@federated_computation.federated_computation(
computation_types.at_server(model_weights_type),
computation_types.at_clients(
computation_types.SequenceType(batch_type)))
def server_eval(server_model_weights: computation_types.FederatedType,
federated_dataset: computation_types.FederatedType):
broadcast_output = broadcast_process.next(
broadcast_process.initialize(), server_model_weights)
client_outputs = intrinsics.federated_map(
client_computation, [broadcast_output.result, federated_dataset])
aggregated_client_outputs = federated_output_computation(
client_outputs)
measurements = intrinsics.federated_zip(
collections.OrderedDict(broadcast=broadcast_output.measurements,
eval=aggregated_client_outputs))
return measurements
return server_eval
def build_training_process(
model_fn: ModelFn,
*, # Callers pass below args by name.
loss_fn: LossFn,
metrics_fn: Optional[MetricsFn] = None,
server_optimizer_fn: OptimizerFn = functools.partial(
tf.keras.optimizers.SGD, 1.0),
client_optimizer_fn: OptimizerFn = functools.partial(
tf.keras.optimizers.SGD, 0.1),
reconstruction_optimizer_fn: OptimizerFn = functools.partial(
tf.keras.optimizers.SGD, 0.1),
dataset_split_fn: Optional[reconstruction_utils.DatasetSplitFn] = None,
client_weighting: Optional[client_weight_lib.ClientWeightType] = None,
broadcast_process: Optional[measured_process_lib.MeasuredProcess] = None,
aggregation_factory: Optional[AggregationFactory] = None,
) -> iterative_process_lib.IterativeProcess:
"""Builds the IterativeProcess for optimization using FedRecon.
Returns a `tff.templates.IterativeProcess` for Federated Reconstruction. On
the client, computation can be divided into two stages: (1) reconstruction of
local variables and (2) training of global variables.
Args:
model_fn: A no-arg function that returns a
`tff.learning.reconstruction.Model`. This method must *not* capture
Tensorflow tensors or variables and use them. must be constructed entirely
from scratch on each invocation, returning the same pre-constructed model
each call will result in an error.
loss_fn: A no-arg function returning a `tf.keras.losses.Loss` to use to
compute local model updates during reconstruction and post-reconstruction
and evaluate the model during training. The final loss metric is the
example-weighted mean loss across batches and across clients. The loss
metric does not include reconstruction batches in the loss.
metrics_fn: A no-arg function returning a list of `tf.keras.metrics.Metric`s
to evaluate the model. Metrics results are computed locally as described
by the metric, and are aggregated across clients as in
`federated_aggregate_keras_metric`. If None, no metrics are applied.
Metrics are not computed on reconstruction batches.
server_optimizer_fn: A `tff.learning.optimizers.Optimizer`, or a no-arg
function that returns a `tf.keras.optimizers.Optimizer` for applying
updates to the global model on the server.
client_optimizer_fn: A `tff.learning.optimizers.Optimizer`, or a no-arg
function that returns a `tf.keras.optimizers.Optimizer` for local client
training after reconstruction.
reconstruction_optimizer_fn: A `tff.learning.optimizers.Optimizer`, or a
no-arg function that returns a `tf.keras.optimizers.Optimizer` used to
reconstruct the local variables, with the global ones frozen, or the first
stage described above.
dataset_split_fn: A `reconstruction_utils.DatasetSplitFn` taking in a single
TF dataset and producing two TF datasets. The first is iterated over
during reconstruction, and the second is iterated over
post-reconstruction. This can be used to preprocess datasets to e.g.
iterate over them for multiple epochs or use disjoint data for
reconstruction and post-reconstruction. If None, split client data in half
for each user, using one half for reconstruction and the other for
evaluation. See `reconstruction_utils.build_dataset_split_fn` for options.
client_weighting: A value of `tff.learning.ClientWeighting` that specifies a
built-in weighting method, or a callable that takes the local metrics of
the model and returns a tensor that provides the weight in the federated
average of model deltas. If None, defaults to weighting by number of
examples.
broadcast_process: A `tff.templates.MeasuredProcess` that broadcasts the
model weights on the server to the clients. It must support the signature
`(input_values@SERVER -> output_values@CLIENT)`. If set to default None,
the server model is broadcast to the clients using the default
`tff.federated_broadcast`.
aggregation_factory: An optional instance of
`tff.aggregators.WeightedAggregationFactory` or
`tff.aggregators.UnweightedAggregationFactory` determining the method of
aggregation to perform. If unspecified, uses a default
`tff.aggregators.MeanFactory` which computes a stateless mean across
clients (weighted depending on `client_weighting`).
Raises:
TypeError: If `broadcast_process` does not have the expected signature.
TypeError: If `aggregation_factory` does not have the expected signature.
ValueError: If `aggregation_factory` is not a
`tff.aggregators.WeightedAggregationFactory` or a
`tff.aggregators.UnweightedAggregationFactory`.
ValueError: If `aggregation_factory` is a
`tff.aggregators.UnweightedAggregationFactory` but `client_weighting` is
not `tff.learning.ClientWeighting.UNIFORM`.
Returns:
A `tff.templates.IterativeProcess`.
"""
with tf.Graph().as_default():
throwaway_model_for_metadata = model_fn()
model_weights_type = type_conversions.type_from_tensors(
reconstruction_utils.get_global_variables(
throwaway_model_for_metadata))
if client_weighting is None:
client_weighting = client_weight_lib.ClientWeighting.NUM_EXAMPLES
if (isinstance(aggregation_factory, factory.UnweightedAggregationFactory)
and client_weighting
is not client_weight_lib.ClientWeighting.UNIFORM):
raise ValueError(
f'Expected `tff.learning.ClientWeighting.UNIFORM` client '
f'weighting with unweighted aggregator, instead got '
f'{client_weighting}')
if broadcast_process is None:
broadcast_process = optimizer_utils.build_stateless_broadcaster(
model_weights_type=model_weights_type)
if not _is_valid_broadcast_process(broadcast_process):
raise TypeError(
'broadcast_process type signature does not conform to expected '
'signature (<state@S, input@S> -> <state@S, result@C, measurements@S>).'
' Got: {t}'.format(t=broadcast_process.next.type_signature))
broadcaster_state_type = (
broadcast_process.initialize.type_signature.result.member)
aggregation_process = _instantiate_aggregation_process(
aggregation_factory, model_weights_type)
aggregator_state_type = (
aggregation_process.initialize.type_signature.result.member)
server_init_tff = _build_server_init_fn(model_fn, server_optimizer_fn,
aggregation_process.initialize,
broadcast_process.initialize)
server_state_type = server_init_tff.type_signature.result.member
server_update_fn = _build_server_update_fn(
model_fn,
server_optimizer_fn,
server_state_type,
server_state_type.model,
aggregator_state_type=aggregator_state_type,
broadcaster_state_type=broadcaster_state_type)
dataset_type = computation_types.SequenceType(
throwaway_model_for_metadata.input_spec)
if dataset_split_fn is None:
dataset_split_fn = reconstruction_utils.build_dataset_split_fn(
split_dataset=True)
client_update_fn = _build_client_update_fn(
model_fn,
loss_fn=loss_fn,
metrics_fn=metrics_fn,
dataset_type=dataset_type,
model_weights_type=server_state_type.model,
client_optimizer_fn=client_optimizer_fn,
reconstruction_optimizer_fn=reconstruction_optimizer_fn,
dataset_split_fn=dataset_split_fn,
client_weighting=client_weighting)
federated_server_state_type = computation_types.at_server(
server_state_type)
federated_dataset_type = computation_types.at_clients(dataset_type)
# Create placeholder metrics to produce a corresponding federated output
# computation.
metrics = []
if metrics_fn is not None:
metrics.extend(metrics_fn())
metrics.append(keras_utils.MeanLossMetric(loss_fn()))
federated_output_computation = (
keras_utils.federated_output_computation_from_metrics(metrics))
run_one_round_tff = _build_run_one_round_fn(
server_update_fn,
client_update_fn,
federated_output_computation,
federated_server_state_type,
federated_dataset_type,
aggregation_process=aggregation_process,
broadcast_process=broadcast_process,
)
process = iterative_process_lib.IterativeProcess(
initialize_fn=server_init_tff, next_fn=run_one_round_tff)
@computations.tf_computation(server_state_type)
def get_model_weights(server_state):
return server_state.model
process.get_model_weights = get_model_weights
return process
def build_federated_averaging_process(
model_fn: Callable[[], model_lib.Model],
client_optimizer_fn: Optional[Callable[
[], tf.keras.optimizers.Optimizer]] = None,
server_optimizer_fn: Callable[
[], tf.keras.optimizers.Optimizer] = DEFAULT_SERVER_OPTIMIZER_FN,
client_weight_fn: Callable[[Any], tf.Tensor] = None,
stateful_delta_aggregate_fn=None,
stateful_model_broadcast_fn=None) -> tff.utils.IterativeProcess:
"""Builds the TFF computations for optimization using federated averaging.
Args:
model_fn: A no-arg function that returns a `tff.learning.Model`.
client_optimizer_fn: An optional no-arg callable that returns a
`tf.keras.Optimizer`
server_optimizer_fn: A no-arg callable that returns a `tf.keras.Optimizer`.
The `apply_gradients` method of this optimizer is used to apply client
updates to the server model. The default creates a
`tf.keras.optimizers.SGD` with a learning rate of 1.0, which simply adds
the average client delta to the server's model.
client_weight_fn: Optional function that takes the output of
`model.report_local_outputs` and returns a tensor that provides the weight
in the federated average of model deltas. If not provided, the default is
the total number of examples processed on device.
stateful_delta_aggregate_fn: A `tff.utils.StatefulAggregateFn` where the
`next_fn` performs a federated aggregation and upates state. That is, it
has TFF type `(state@SERVER, value@CLIENTS, weights@CLIENTS) ->
(state@SERVER, aggregate@SERVER)`, where the `value` type is
`tff.learning.framework.ModelWeights.trainable` corresponding to the
object returned by `model_fn`. By default performs arithmetic mean
aggregation, weighted by `client_weight_fn`.
stateful_model_broadcast_fn: A `tff.utils.StatefulBroadcastFn` where the
`next_fn` performs a federated broadcast and upates state. That is, it has
TFF type `(state@SERVER, value@SERVER) -> (state@SERVER, value@CLIENTS)`,
where the `value` type is `tff.learning.framework.ModelWeights`
corresponding to the object returned by `model_fn`. By default performs
identity broadcast.
Returns:
A `tff.utils.IterativeProcess`.
"""
if client_optimizer_fn is None:
warnings.warn('tff.learning.build_federated_averaging_process will start '
'requiring a new argument \'client_optimizer_fn\'. Specify '
'the local client optimizer here rather than building a '
'ttf.learning.TrainableModel')
else:
# Validate parameters and surfacing errors early requires building a
# throwaway model here.
with tf.Graph().as_default():
model = model_fn()
if isinstance(model, model_lib.TrainableModel):
raise TypeError('model_fn parameter should be a callable that produces '
'tff.learning.Model, not the deprecated '
'tff.learning.TrainableModel')
def client_fed_avg(model_fn):
if client_optimizer_fn is None:
return _DeprecatedClientFedAvg(model_fn(), client_weight_fn)
elif callable(client_optimizer_fn):
return _ClientFedAvg(model_fn(), client_optimizer_fn(), client_weight_fn)
else:
raise TypeError(f'client_optimizer_fn parameter of '
'tff.learning.build_federated_averaging_process must be '
'a callable. Received a {type(client_optimizer_fn)}')
if stateful_delta_aggregate_fn is None:
stateful_delta_aggregate_fn = optimizer_utils.build_stateless_mean()
else:
py_typecheck.check_type(stateful_delta_aggregate_fn,
tff.utils.StatefulAggregateFn)
if stateful_model_broadcast_fn is None:
stateful_model_broadcast_fn = optimizer_utils.build_stateless_broadcaster()
else:
py_typecheck.check_type(stateful_model_broadcast_fn,
tff.utils.StatefulBroadcastFn)
return optimizer_utils.build_model_delta_optimizer_process(
model_fn, client_fed_avg, server_optimizer_fn,
stateful_delta_aggregate_fn, stateful_model_broadcast_fn)
