/
measurements.py
181 lines (155 loc) · 7.78 KB
/
measurements.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
# Copyright 2021, The TensorFlow Federated Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Aggregation factory for adding custom measurements."""
from collections.abc import Callable
import inspect
import typing
from typing import Any, Optional
from tensorflow_federated.python.aggregators import factory
from tensorflow_federated.python.common_libs import py_typecheck
from tensorflow_federated.python.core.environments.tensorflow_frontend import tensorflow_computation
from tensorflow_federated.python.core.impl.federated_context import federated_computation
from tensorflow_federated.python.core.impl.federated_context import intrinsics
from tensorflow_federated.python.core.impl.types import computation_types
from tensorflow_federated.python.core.impl.types import placements
from tensorflow_federated.python.core.templates import aggregation_process
from tensorflow_federated.python.core.templates import measured_process
def add_measurements(
inner_agg_factory: factory.AggregationFactory,
*,
client_measurement_fn: Optional[Callable[..., dict[str, Any]]] = None,
server_measurement_fn: Optional[Callable[..., dict[str, Any]]] = None,
) -> factory.AggregationFactory:
"""Wraps `AggregationFactory` to report additional measurements.
The function `client_measurement_fn` should be a Python callable that will be
called as `client_measurement_fn(value)` or `client_measurement_fn(value,
weight)` depending on whether `inner_agg_factory` is weighted or unweighted.
It must be traceable by TFF and expect `tff.Value` objects placed at `CLIENTS`
as inputs, and return a `collections.OrderedDict` mapping string names to
tensor values placed at `SERVER`, which will be added to the measurement dict
produced by the `inner_agg_factory`.
Similarly, `server_measurement_fn` should be a Python callable that will be
called as `server_measurement_fn(result)` where `result` is the result (on
server) of the inner aggregation.
One or both of `client_measurement_fn` and `server_measurement_fn` must be
specified.
Args:
inner_agg_factory: The factory to wrap and add measurements.
client_measurement_fn: A Python callable that will be called on `value`
(and/or `weight`) provided to the `next` function to compute additional
measurements of the client values/weights.
server_measurement_fn: A Python callable that will be called on the `result`
of aggregation at server to compute additional measurements of the result.
Returns:
An `AggregationFactory` that reports additional measurements.
"""
type_args = typing.get_args(factory.AggregationFactory)
py_typecheck.check_type(inner_agg_factory, type_args)
if client_measurement_fn is None and server_measurement_fn is None:
raise ValueError(
'Must specify one or both of `client_measurement_fn` or '
'`server_measurement_fn`.'
)
if client_measurement_fn is not None:
if isinstance(inner_agg_factory, factory.UnweightedAggregationFactory):
if len(inspect.signature(client_measurement_fn).parameters) != 1:
raise ValueError(
'`client_measurement_fn` must take a single parameter if '
'`inner_agg_factory` is unweighted.'
)
elif isinstance(inner_agg_factory, factory.WeightedAggregationFactory):
if len(inspect.signature(client_measurement_fn).parameters) != 2:
raise ValueError(
'`client_measurement_fn` must take a two parameters if '
'`inner_agg_factory` is weighted.'
)
if server_measurement_fn is not None:
if len(inspect.signature(server_measurement_fn).parameters) != 1:
raise ValueError('`server_measurement_fn` must take a single parameter.')
@tensorflow_computation.tf_computation()
def dict_update(orig_dict, new_values):
if not orig_dict:
return new_values
orig_dict.update(new_values)
return orig_dict
if isinstance(inner_agg_factory, factory.WeightedAggregationFactory):
class WeightedWrappedFactory(factory.WeightedAggregationFactory):
"""Wrapper for `WeightedAggregationFactory` that adds new measurements."""
def create(
self, value_type: factory.ValueType, weight_type: factory.ValueType
) -> aggregation_process.AggregationProcess:
type_args = typing.get_args(factory.ValueType)
py_typecheck.check_type(value_type, type_args)
py_typecheck.check_type(weight_type, type_args)
inner_agg_process = inner_agg_factory.create(value_type, weight_type)
init_fn = inner_agg_process.initialize
@federated_computation.federated_computation(
init_fn.type_signature.result,
computation_types.FederatedType(value_type, placements.CLIENTS),
computation_types.FederatedType(weight_type, placements.CLIENTS),
)
def next_fn(state, value, weight):
inner_agg_output = inner_agg_process.next(state, value, weight)
measurements = inner_agg_output.measurements
if client_measurement_fn:
client_measurements = client_measurement_fn(value, weight)
measurements = intrinsics.federated_map(
dict_update, (measurements, client_measurements)
)
if server_measurement_fn:
server_measurements = server_measurement_fn(inner_agg_output.result)
measurements = intrinsics.federated_map(
dict_update, (measurements, server_measurements)
)
return measured_process.MeasuredProcessOutput(
state=inner_agg_output.state,
result=inner_agg_output.result,
measurements=measurements,
)
return aggregation_process.AggregationProcess(init_fn, next_fn)
return WeightedWrappedFactory()
else:
class UnweightedWrappedFactory(factory.UnweightedAggregationFactory):
"""Wrapper for `UnweightedAggregationFactory` that adds new measurements."""
def create(
self, value_type: factory.ValueType
) -> aggregation_process.AggregationProcess:
type_args = typing.get_args(factory.ValueType)
py_typecheck.check_type(value_type, type_args)
inner_agg_process = inner_agg_factory.create(value_type)
init_fn = inner_agg_process.initialize
@federated_computation.federated_computation(
init_fn.type_signature.result,
computation_types.FederatedType(value_type, placements.CLIENTS),
)
def next_fn(state, value):
inner_agg_output = inner_agg_process.next(state, value)
measurements = inner_agg_output.measurements
if client_measurement_fn:
client_measurements = client_measurement_fn(value)
measurements = intrinsics.federated_map(
dict_update, (measurements, client_measurements)
)
if server_measurement_fn:
server_measurements = server_measurement_fn(inner_agg_output.result)
measurements = intrinsics.federated_map(
dict_update, (measurements, server_measurements)
)
return measured_process.MeasuredProcessOutput(
state=inner_agg_output.state,
result=inner_agg_output.result,
measurements=measurements,
)
return aggregation_process.AggregationProcess(init_fn, next_fn)
return UnweightedWrappedFactory()