def test_federated_aggregate_with_client_int(self):
# The representation used during the aggregation process will be a named
# tuple with 2 elements - the integer 'total' that represents the sum of
# elements encountered, and the integer element 'count'.
# pylint: disable=invalid-name
Accumulator = collections.namedtuple('Accumulator', 'total count')
# pylint: enable=invalid-name
# The operator to use during the first stage simply adds an element to the
# total and updates the count.
@computations.tf_computation
def accumulate(accu, elem):
return Accumulator(accu.total + elem, accu.count + 1)
# The operator to use during the second stage simply adds total and count.
@computations.tf_computation
def merge(x, y):
return Accumulator(x.total + y.total, x.count + y.count)
# The operator to use during the final stage simply computes the ratio.
@computations.tf_computation
def report(accu):
return tf.cast(accu.total, tf.float32) / tf.cast(
accu.count, tf.float32)
x = _mock_data_of_type(computation_types.at_clients(tf.int32))
val = intrinsics.federated_aggregate(x, Accumulator(0, 0), accumulate,
merge, report)
self.assert_value(val, 'float32@SERVER')
def _federated_reduce_with_func(value, tf_func, zeros):
"""Applies to `tf_func` to accumulated `value`s.
This utility provides a generic aggregation for accumulating a value and
applying a simple aggregation (like minimum or maximum aggregations).
Args:
value: A `tff.Value` placed on the `tff.CLIENTS`.
tf_func: A function to be applied to the accumulated values. Must be a
binary operation where both parameters are of type `U` and the return type
is also `U`.
zeros: The zero of the same type as `value` in the algebra of reduction
operators.
Returns:
A representation on the `tff.SERVER` of the result of aggregating `value`.
"""
value_type = value.type_signature.member
@tensorflow_computation.tf_computation(value_type, value_type)
def accumulate(current, value):
return tf.nest.map_structure(tf_func, current, value)
@tensorflow_computation.tf_computation(value_type)
def report(value):
return value
return intrinsics.federated_aggregate(value, zeros, accumulate, accumulate,
report)
def next_fn(state, value):
encode_params, decode_before_sum_params, decode_after_sum_params = (
intrinsics.federated_map(get_params_fn, state))
encode_params = intrinsics.federated_broadcast(encode_params)
decode_before_sum_params = intrinsics.federated_broadcast(
decode_before_sum_params)
encoded_values = intrinsics.federated_map(
encode_fn, [value, encode_params, decode_before_sum_params])
aggregated_values = intrinsics.federated_aggregate(
encoded_values, zero_fn(), accumulate_fn, merge_fn, report_fn)
decoded_values = intrinsics.federated_map(
decode_after_sum_fn,
[aggregated_values.values, decode_after_sum_params])
updated_state = intrinsics.federated_map(
update_state_fn, [state, aggregated_values.state_update_tensors])
empty_metrics = intrinsics.federated_value((), placements.SERVER)
return measured_process.MeasuredProcessOutput(
state=updated_state,
result=decoded_values,
measurements=empty_metrics)
def test_federated_aggregate_with_unknown_dimension(self):
Accumulator = collections.namedtuple('Accumulator', ['samples']) # pylint: disable=invalid-name
accumulator_type = computation_types.StructType(
Accumulator(samples=computation_types.TensorType(dtype=tf.int32,
shape=[None])))
@computations.tf_computation()
def build_empty_accumulator():
return Accumulator(samples=tf.zeros(shape=[0], dtype=tf.int32))
# The operator to use during the first stage simply adds an element to the
# tensor, increasing its size.
@computations.tf_computation([accumulator_type, tf.int32])
def accumulate(accu, elem):
return Accumulator(samples=tf.concat(
[accu.samples, tf.expand_dims(elem, axis=0)], axis=0))
# The operator to use during the second stage simply adds total and count.
@computations.tf_computation([accumulator_type, accumulator_type])
def merge(x, y):
return Accumulator(
samples=tf.concat([x.samples, y.samples], axis=0))
# The operator to use during the final stage simply computes the ratio.
@computations.tf_computation(accumulator_type)
def report(accu):
return accu
x = _mock_data_of_type(computation_types.at_clients(tf.int32))
val = intrinsics.federated_aggregate(x, build_empty_accumulator(),
accumulate, merge, report)
self.assert_value(val, '<samples=int32[?]>@SERVER')
def up_to_merge_computation():
federated_aggregate_args = before_agg_callable(
)['federated_aggregate_param']
value_to_aggregate = federated_aggregate_args[0]
zero = zero_comp()
return intrinsics.federated_aggregate(value_to_aggregate, zero,
accumulate_comp, merge_comp,
identity_report)
def comp():
value = intrinsics.federated_value(
collections.OrderedDict([('a', (10, 2.0))]),
placements.CLIENTS)
zero = collections.OrderedDict([('a', (0, 0.0))])
return intrinsics.federated_aggregate(value, zero, add_test_type,
add_test_type,
add_five_and_three)
def test_federated_aggregate_with_federated_zero_fails(self):
x = _mock_data_of_type(computation_types.at_clients(tf.int32))
zero = intrinsics.federated_value(0, placements.SERVER)
accumulate = _create_computation_add()
# The operator to use during the second stage simply adds total and count.
merge = _create_computation_add()
# The operator to use during the final stage simply computes the ratio.
report_type = computation_types.TensorType(tf.int32)
report_proto, _ = tensorflow_computation_factory.create_identity(
report_type)
report = computation_impl.ConcreteComputation(
report_proto, context_stack_impl.context_stack)
with self.assertRaisesRegex(
TypeError, 'Expected `zero` to be assignable to type int32, '
'but was of incompatible type int32@SERVER'):
intrinsics.federated_aggregate(x, zero, accumulate, merge, report)
def test_federated_aggregate_with_unknown_dimension(self):
Accumulator = collections.namedtuple('Accumulator', ['samples']) # pylint: disable=invalid-name
accumulator_type = computation_types.to_type(
Accumulator(samples=computation_types.TensorType(dtype=tf.int32,
shape=[None])))
x = _mock_data_of_type(computation_types.at_clients(tf.int32))
def initialize_fn():
return Accumulator(samples=tf.zeros(shape=[0], dtype=tf.int32))
initialize_proto, _ = tensorflow_computation_factory.create_computation_for_py_fn(
initialize_fn, None)
initialize = computation_impl.ConcreteComputation(
initialize_proto, context_stack_impl.context_stack)
zero = initialize()
# The operator to use during the first stage simply adds an element to the
# tensor, increasing its size.
def _accumulate(arg):
return Accumulator(samples=tf.concat(
[arg[0].samples,
tf.expand_dims(arg[1], axis=0)], axis=0))
accumulate_type = computation_types.StructType(
[accumulator_type, tf.int32])
accumulate_proto, _ = tensorflow_computation_factory.create_computation_for_py_fn(
_accumulate, accumulate_type)
accumulate = computation_impl.ConcreteComputation(
accumulate_proto, context_stack_impl.context_stack)
# The operator to use during the second stage simply adds total and count.
def _merge(arg):
return Accumulator(
samples=tf.concat([arg[0].samples, arg[1].samples], axis=0))
merge_type = computation_types.StructType(
[accumulator_type, accumulator_type])
merge_proto, _ = tensorflow_computation_factory.create_computation_for_py_fn(
_merge, merge_type)
merge = computation_impl.ConcreteComputation(
merge_proto, context_stack_impl.context_stack)
# The operator to use during the final stage simply computes the ratio.
report_proto, _ = tensorflow_computation_factory.create_identity(
accumulator_type)
report = computation_impl.ConcreteComputation(
report_proto, context_stack_impl.context_stack)
value = intrinsics.federated_aggregate(x, zero, accumulate, merge,
report)
self.assert_value(value, '<samples=int32[?]>@SERVER')
def test_federated_aggregate_with_federated_zero_fails(self):
zero = intrinsics.federated_value(0, placements.SERVER)
@computations.tf_computation([tf.int32, tf.int32])
def accumulate(accu, elem):
return accu + elem
# The operator to use during the second stage simply adds total and count.
@computations.tf_computation([tf.int32, tf.int32])
def merge(x, y):
return x + y
# The operator to use during the final stage simply computes the ratio.
@computations.tf_computation(tf.int32)
def report(accu):
return accu
x = _mock_data_of_type(computation_types.at_clients(tf.int32))
with self.assertRaisesRegex(
TypeError, 'Expected `zero` to be assignable to type int32, '
'but was of incompatible type int32@SERVER'):
intrinsics.federated_aggregate(x, zero, accumulate, merge, report)
def test_federated_aggregate_with_client_int(self):
# The representation used during the aggregation process will be a named
# tuple with 2 elements - the integer 'total' that represents the sum of
# elements encountered, and the integer element 'count'.
Accumulator = collections.namedtuple('Accumulator', 'total count') # pylint: disable=invalid-name
accumulator_type = computation_types.to_type(
Accumulator(total=computation_types.TensorType(dtype=tf.int32),
count=computation_types.TensorType(dtype=tf.int32)))
x = _mock_data_of_type(computation_types.at_clients(tf.int32))
zero = Accumulator(0, 0)
# The operator to use during the first stage simply adds an element to the
# total and updates the count.
def _accumulate(arg):
return Accumulator(arg[0].total + arg[1], arg[0].count + 1)
accumulate_type = computation_types.StructType(
[accumulator_type, tf.int32])
accumulate_proto, _ = tensorflow_computation_factory.create_computation_for_py_fn(
_accumulate, accumulate_type)
accumulate = computation_impl.ConcreteComputation(
accumulate_proto, context_stack_impl.context_stack)
# The operator to use during the second stage simply adds total and count.
def _merge(arg):
return Accumulator(arg[0].total + arg[1].total,
arg[0].count + arg[1].count)
merge_type = computation_types.StructType(
[accumulator_type, accumulator_type])
merge_proto, _ = tensorflow_computation_factory.create_computation_for_py_fn(
_merge, merge_type)
merge = computation_impl.ConcreteComputation(
merge_proto, context_stack_impl.context_stack)
# The operator to use during the final stage simply computes the ratio.
def _report(arg):
return tf.cast(arg.total, tf.float32) / tf.cast(
arg.count, tf.float32)
report_proto, _ = tensorflow_computation_factory.create_computation_for_py_fn(
_report, accumulator_type)
report = computation_impl.ConcreteComputation(
report_proto, context_stack_impl.context_stack)
value = intrinsics.federated_aggregate(x, zero, accumulate, merge,
report)
self.assert_value(value, 'float32@SERVER')
def next_fn(unused_state, value):
# Empty tuple is the `None` of TFF.
empty_tuple = intrinsics.federated_value((), placements.SERVER)
initial_reservoir = _build_initial_sample_reservoir(value_type)
sample_value = _build_sample_value_computation(
value_type, self._sample_size)
merge_samples = _build_merge_samples_computation(
value_type, self._sample_size)
finalize_sample = _build_finalize_sample_computation(value_type)
samples = intrinsics.federated_aggregate(value,
zero=initial_reservoir,
accumulate=sample_value,
merge=merge_samples,
report=finalize_sample)
return measured_process.MeasuredProcessOutput(
state=empty_tuple, result=samples, measurements=empty_tuple)
def next_computation(arg):
"""The logic of a single MapReduce processing round."""
s1 = arg[0]
c1 = arg[1]
s2 = intrinsics.federated_map(mrf.prepare, s1)
c2 = intrinsics.federated_broadcast(s2)
c3 = intrinsics.federated_zip([c1, c2])
c4 = intrinsics.federated_map(mrf.work, c3)
c5 = c4[0]
c6 = c4[1]
s3 = intrinsics.federated_aggregate(c5, mrf.zero(), mrf.accumulate,
mrf.merge, mrf.report)
s4 = intrinsics.federated_secure_sum_bitwidth(c6, mrf.bitwidth())
s5 = intrinsics.federated_zip([s3, s4])
s6 = intrinsics.federated_zip([s1, s5])
s7 = intrinsics.federated_map(mrf.update, s6)
s8 = s7[0]
s9 = s7[1]
return s8, s9
def test_infers_accumulate_return_as_merge_arg_merge_return_as_report_arg(
self):
type_spec = computation_types.TensorType(dtype=tf.int64, shape=[None])
x = _mock_data_of_type(computation_types.at_clients(tf.int64))
def initialize_fn():
return tf.constant([], dtype=tf.int64, shape=[0])
initialize_proto, _ = tensorflow_computation_factory.create_computation_for_py_fn(
initialize_fn, None)
initialize = computation_impl.ConcreteComputation(
initialize_proto, context_stack_impl.context_stack)
zero = initialize()
def _accumulate(arg):
return tf.concat([arg[0], [arg[1]]], 0)
accumulate_type = computation_types.StructType([type_spec, tf.int64])
accumulate_proto, _ = tensorflow_computation_factory.create_computation_for_py_fn(
_accumulate, accumulate_type)
accumulate = computation_impl.ConcreteComputation(
accumulate_proto, context_stack_impl.context_stack)
def _merge(arg):
return tf.concat([arg[0], arg[1]], 0)
merge_type = computation_types.StructType([type_spec, type_spec])
merge_proto, _ = tensorflow_computation_factory.create_computation_for_py_fn(
_merge, merge_type)
merge = computation_impl.ConcreteComputation(
merge_proto, context_stack_impl.context_stack)
report_proto, _ = tensorflow_computation_factory.create_identity(
type_spec)
report = computation_impl.ConcreteComputation(
report_proto, context_stack_impl.context_stack)
value = intrinsics.federated_aggregate(x, zero, accumulate, merge,
report)
self.assert_value(value, 'int64[?]@SERVER')
def next_computation(arg):
"""The logic of a single MapReduce processing round."""
server_state, client_data = arg
broadcast_input = intrinsics.federated_map(mrf.prepare, server_state)
broadcast_result = intrinsics.federated_broadcast(broadcast_input)
work_arg = intrinsics.federated_zip([client_data, broadcast_result])
(aggregate_input, secure_sum_bitwidth_input, secure_sum_input,
secure_modular_sum_input) = intrinsics.federated_map(
mrf.work, work_arg)
aggregate_result = intrinsics.federated_aggregate(
aggregate_input, mrf.zero(), mrf.accumulate, mrf.merge, mrf.report)
secure_sum_bitwidth_result = intrinsics.federated_secure_sum_bitwidth(
secure_sum_bitwidth_input, mrf.secure_sum_bitwidth())
secure_sum_result = intrinsics.federated_secure_sum(
secure_sum_input, mrf.secure_sum_max_input())
secure_modular_sum_result = intrinsics.federated_secure_modular_sum(
secure_modular_sum_input, mrf.secure_modular_sum_modulus())
update_arg = intrinsics.federated_zip(
(server_state, (aggregate_result, secure_sum_bitwidth_result,
secure_sum_result, secure_modular_sum_result)))
updated_server_state, server_output = intrinsics.federated_map(
mrf.update, update_arg)
return updated_server_state, server_output
def federated_aggregate_keras_metric(
metrics: Union[tf.keras.metrics.Metric,
Sequence[tf.keras.metrics.Metric]], federated_values):
"""Aggregates variables a keras metric placed at CLIENTS to SERVER.
Args:
metrics: A single `tf.keras.metrics.Metric` or a `Sequence` of metrics . The
order must match the order of variables in `federated_values`.
federated_values: A single federated value, or a `Sequence` of federated
values. The values must all have `tff.CLIENTS` placement. If value is a
`Sequence` type, it must match the order of the sequence in `metrics.
Returns:
The result of performing a federated sum on federated_values, then assigning
the aggregated values into the variables of the corresponding
`tf.keras.metrics.Metric` and calling `tf.keras.metrics.Metric.result`. The
resulting structure has `tff.SERVER` placement.
"""
member_types = tf.nest.map_structure(lambda t: t.type_signature.member,
federated_values)
@computations.tf_computation
def zeros_fn():
# `member_type` is a (potentially nested) `tff.StructType`, which is an
# `structure.Struct`.
return structure.map_structure(
lambda v: tf.zeros(v.shape, dtype=v.dtype), member_types)
zeros = zeros_fn()
@computations.tf_computation(member_types, member_types)
def accumulate(accumulators, variables):
return tf.nest.map_structure(tf.add, accumulators, variables)
@computations.tf_computation(member_types, member_types)
def merge(a, b):
return tf.nest.map_structure(tf.add, a, b)
@computations.tf_computation(member_types)
def report(accumulators):
"""Insert `accumulators` back into the keras metric to obtain result."""
def finalize_metric(metric: tf.keras.metrics.Metric, values):
# Note: the following call requires that `type(metric)` have a no argument
# __init__ method, which will restrict the types of metrics that can be
# used. This is somewhat limiting, but the pattern to use default
# arguments and export the values in `get_config()` (see
# `tf.keras.metrics.TopKCategoricalAccuracy`) works well.
#
# If type(metric) is subclass of another tf.keras.metric arguments passed
# to __init__ must include arguments expected by the superclass and
# specified in superclass get_config().
keras_metric = None
try:
# This is some trickery to reconstruct a metric object in the current
# scope, so that the `tf.Variable`s get created when we desire.
keras_metric = type(metric).from_config(metric.get_config())
except TypeError as e:
# Re-raise the error with a more helpful message, but the previous stack
# trace.
raise TypeError(
'Caught exception trying to call `{t}.from_config()` with '
'config {c}. Confirm that {t}.__init__() has an argument for '
'each member of the config.\nException: {e}'.format(
t=type(metric), c=metric.get_config(), e=e))
assignments = []
for v, a in zip(keras_metric.variables, values):
assignments.append(v.assign(a))
with tf.control_dependencies(assignments):
return keras_metric.result()
if isinstance(metrics, tf.keras.metrics.Metric):
# Only a single metric to aggregate.
return finalize_metric(metrics, accumulators)
else:
# Otherwise map over all the metrics.
return collections.OrderedDict([
(name, finalize_metric(metric, values))
for metric, (name,
values) in zip(metrics, accumulators.items())
])
return intrinsics.federated_aggregate(federated_values, zeros, accumulate,
merge, report)
def comp():
value = intrinsics.federated_value(10, placements.CLIENTS)
return intrinsics.federated_aggregate(value, 0, add_int, add_int,
add_five)
def federated_sample(value, max_num_samples=100):
"""Aggregation to produce uniform sample of at most `max_num_samples` values.
Each client value is assigned a random number when it is examined during each
accumulation. Each accumulate and merge only keeps the top N values based
on the random number. Report drops the random numbers and only returns the
at most N values sampled from the accumulated client values using standard
reservoir sampling (https://en.wikipedia.org/wiki/Reservoir_sampling), where
N is user provided `max_num_samples`.
Args:
value: A `tff.Value` placed on the `tff.CLIENTS`.
max_num_samples: The maximum number of samples to collect from client
values. If fewer clients than the defined max sample size participated in
the round of computation, the actual number of samples will equal the
number of clients in the round.
Returns:
At most `max_num_samples` samples of the value from the `tff.CLIENTS`.
"""
_validate_value_on_clients(value)
member_type = value.type_signature.member
accumulator_type = _get_accumulator_type(member_type)
zeros = _zeros_for_sample(member_type)
def apply_sampling(accumulators, rands):
size = tf.shape(rands)[0]
k = tf.minimum(size, max_num_samples)
indices = tf.math.top_k(rands, k=k).indices
return tf.nest.map_structure(lambda v: tf.gather(v, indices),
accumulators), tf.gather(rands, indices)
def concat_expand_dims(a, b):
b = tf.expand_dims(b, axis=0)
return tf.concat([a, b], axis=0)
@tensorflow_computation.tf_computation(accumulator_type,
value.type_signature.member)
def accumulate(current, value):
"""Accumulates samples through concatenation."""
rands = concat_expand_dims(current.rands, tf.random.uniform(shape=()))
accumulators = tf.nest.map_structure(concat_expand_dims,
current.accumulators, value)
accumulators, rands = apply_sampling(accumulators, rands)
return _Samples(accumulators, rands)
@tensorflow_computation.tf_computation(accumulator_type, accumulator_type)
def merge(a, b):
"""Merges accumulators through concatenation."""
def zero_axis_concat(a, b):
return tf.concat([a, b], axis=0)
samples = tf.nest.map_structure(zero_axis_concat, a, b)
accumulators, rands = apply_sampling(samples.accumulators, samples.rands)
return _Samples(accumulators, rands)
@tensorflow_computation.tf_computation(accumulator_type)
def report(value):
return value.accumulators
return intrinsics.federated_aggregate(value, zeros, accumulate, merge, report)
def federated_aggregate_keras_metric(
metrics: Union[tf.keras.metrics.Metric, Sequence[tf.keras.metrics.Metric],
Callable[[], tf.keras.metrics.Metric],
Sequence[Callable[[], tf.keras.metrics.Metric]]],
federated_values):
"""Aggregates variables a keras metric placed at CLIENTS to SERVER.
Args:
metrics: A single or a `Sequence` of `tf.keras.metrics.Metric` objects, or a
single or a `Sequence` of no-arg callables that each constructs a
`tf.keras.metrics.Metric`. The order must match the order of variables in
`federated_values`.
federated_values: A single federated value, or a `Sequence` of federated
values. The values must all have `tff.CLIENTS` placement. If value is a
`Sequence` type, it must match the order of the sequence in `metrics.
Returns:
The result of performing a federated sum on federated_values, then assigning
the aggregated values into the variables of the corresponding
`tf.keras.metrics.Metric` and calling `tf.keras.metrics.Metric.result`. The
resulting structure has `tff.SERVER` placement.
"""
member_types = tf.nest.map_structure(lambda t: t.type_signature.member,
federated_values)
@tensorflow_computation.tf_computation
def zeros_fn():
return type_conversions.structure_from_tensor_type_tree(
lambda t: tf.zeros(shape=t.shape, dtype=t.dtype), member_types)
zeros = zeros_fn()
@tensorflow_computation.tf_computation(member_types, member_types)
def accumulate(accumulators, variables):
return tf.nest.map_structure(tf.add, accumulators, variables)
@tensorflow_computation.tf_computation(member_types, member_types)
def merge(a, b):
return tf.nest.map_structure(tf.add, a, b)
@tensorflow_computation.tf_computation(member_types)
def report(accumulators):
"""Insert `accumulators` back into the keras metric to obtain result."""
def finalize_metric(metric: Union[tf.keras.metrics.Metric,
Callable[[], tf.keras.metrics.Metric]],
values):
# Note: if the input metric is an instance of `tf.keras.metrics.Metric`,
# the following call requires that `type(metric)` have a no argument
# __init__ method, which will restrict the types of metrics that can be
# used. This is somewhat limiting, but the pattern to use default
# arguments and export the values in `get_config()` (see
# `tf.keras.metrics.TopKCategoricalAccuracy`) works well.
#
# If type(metric) is subclass of another tf.keras.metric arguments passed
# to __init__ must include arguments expected by the superclass and
# specified in superclass get_config().
keras_metric = finalizer.create_keras_metric(metric)
assignments = []
for v, a in zip(keras_metric.variables, values):
assignments.append(v.assign(a))
with tf.control_dependencies(assignments):
return keras_metric.result()
if isinstance(metrics, tf.keras.metrics.Metric):
# Only a single metric to aggregate.
return finalize_metric(metrics, accumulators)
else:
# Otherwise map over all the metrics.
return collections.OrderedDict([
(name, finalize_metric(metric, values))
for metric, (name, values) in zip(metrics, accumulators.items())
])
return intrinsics.federated_aggregate(federated_values, zeros, accumulate,
merge, report)
