def test_construction_calls_model_fn(self):
# Assert that the the process building does not call `model_fn` too many
# times. `model_fn` can potentially be expensive (loading weights,
# processing, etc).
mock_model_fn = mock.Mock(side_effect=model_examples.LinearRegression)
federated_averaging.build_federated_averaging_process(
model_fn=mock_model_fn,
client_optimizer_fn=tf.keras.optimizers.SGD)
# TODO(b/186451541): reduce the number of calls to model_fn.
self.assertEqual(mock_model_fn.call_count, 4)
def test_fails_stateful_aggregate_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_delta_aggregate_fn=tff.utils.StatefulAggregateFn(
initialize_fn=lambda: (),
next_fn=lambda state, value, weight=None: # pylint: disable=g-long-lambda
(state, tff.federated_mean(value, weight))),
aggregation_process=optimizer_utils.build_stateless_mean(
model_delta_type=model_weights_type.trainable))
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_aggregation_process_deprecation(self):
aggregation_process = mean.MeanFactory().create(
computation_types.to_type([(tf.float32, (2, 1)), tf.float32]),
computation_types.TensorType(tf.float32))
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter('always')
federated_averaging.build_federated_averaging_process(
model_fn=model_examples.LinearRegression,
client_optimizer_fn=lambda: tf.keras.optimizers.SGD(0.1),
aggregation_process=aggregation_process)
self.assertNotEmpty(w)
self.assertEqual(w[0].category, DeprecationWarning)
self.assertRegex(str(w[0].message),
'aggregation_process .* is deprecated')
def test_orchestration_execute_from_keras(self, build_keras_model_fn):
dummy_batch = collections.OrderedDict(
x=np.zeros([1, 2], np.float32), y=np.zeros([1, 1], np.float32))
def model_fn():
keras_model = build_keras_model_fn(feature_dims=2)
return keras_utils.from_keras_model(
keras_model,
dummy_batch,
loss=tf.keras.losses.MeanSquaredError(),
metrics=[])
iterative_process = federated_averaging.build_federated_averaging_process(
model_fn=model_fn,
client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.01))
ds = tf.data.Dataset.from_tensor_slices(
collections.OrderedDict(
x=[[1.0, 2.0], [3.0, 4.0]],
y=[[5.0], [6.0]],
)).batch(2)
federated_ds = [ds] * 3
server_state = iterative_process.initialize()
prev_loss = np.inf
for _ in range(3):
server_state, metrics = iterative_process.next(server_state, federated_ds)
self.assertLess(metrics.loss, prev_loss)
prev_loss = metrics.loss
def test_orchestration_execute_from_keras(self, build_keras_model_fn):
dummy_batch = collections.OrderedDict([
('x', np.zeros([1, 2], np.float32)),
('y', np.zeros([1, 1], np.float32)),
])
def model_fn():
keras_model = build_keras_model_fn(feature_dims=2)
keras_model.compile(
optimizer=tf.keras.optimizers.SGD(learning_rate=0.01),
loss=tf.keras.losses.MeanSquaredError(),
metrics=[])
return keras_utils.from_compiled_keras_model(keras_model, dummy_batch)
iterative_process = federated_averaging.build_federated_averaging_process(
model_fn=model_fn)
ds = tf.data.Dataset.from_tensor_slices({
'x': [[1., 2.], [3., 4.]],
'y': [[5.], [6.]]
}).batch(2)
federated_ds = [ds] * 3
server_state = iterative_process.initialize()
prev_loss = np.inf
for _ in range(3):
server_state, metrics = iterative_process.next(server_state, federated_ds)
self.assertLess(metrics.loss, prev_loss)
prev_loss = metrics.loss
async def fed_avg():
await load_data()
iterative_process = federated_averaging.build_federated_averaging_process(
model_fn=model_examples.LinearRegression,
client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.1))
model = iterative_process.initialize()
self.assertIsInstance(
iterative_process.get_model_weights(model), model_utils.ModelWeights)
self.assertAllClose(model.model.trainable,
iterative_process.get_model_weights(model).trainable)
for _ in range(num_clients):
model, _ = iterative_process.next(model, datasets)
# self.assertIsInstance(
# iterative_process.get_model_weights(model), model_utils.ModelWeights)
# self.assertAllClose(model.model.trainable,
# iterative_process.get_model_weights(model).trainable)
model.save_weights('consolidated.h5', save_format = 'h5')
update_model(model)
def test_orchestration_execute_from_keras_with_lookup(self):
self.skipTest('https://github.com/tensorflow/federated/issues/783')
def model_fn():
dummy_batch = collections.OrderedDict([
('x', tf.constant([['R']], tf.string)),
('y', tf.zeros([1, 1], tf.float32)),
])
keras_model = model_examples.build_lookup_table_keras_model()
return keras_utils.from_compiled_keras_model(
keras_model,
dummy_batch,
loss=tf.keras.losses.MeanSquaredError(),
metrics=[])
iterative_process = federated_averaging.build_federated_averaging_process(
model_fn=model_fn,
client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=
0.1))
ds = tf.data.Dataset.from_tensor_slices(
collections.OrderedDict([
('x', [['R'], ['G'], ['B']]),
('y', [[1.0], [2.0], [3.0]]),
])).batch(2)
federated_ds = [ds] * 3
server_state = iterative_process.initialize()
prev_loss = np.inf
for _ in range(3):
server_state, metrics = iterative_process.next(
server_state, federated_ds)
self.assertLess(metrics.loss, prev_loss)
prev_loss = metrics.loss
def test_orchestration_execute_from_keras_with_lookup(self):
ds = tf.data.Dataset.from_tensor_slices(
collections.OrderedDict(x=[['R'], ['G'], ['B']],
y=[[1.0], [2.0], [3.0]])).batch(2)
def model_fn():
keras_model = model_examples.build_lookup_table_keras_model()
return keras_utils.from_keras_model(
keras_model,
loss=tf.keras.losses.MeanSquaredError(),
input_spec=ds.element_spec,
metrics=[])
iterative_process = federated_averaging.build_federated_averaging_process(
model_fn=model_fn,
client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=
0.1))
federated_ds = [ds] * 3
server_state = iterative_process.initialize()
prev_loss = np.inf
for _ in range(3):
server_state, metrics = iterative_process.next(
server_state, federated_ds)
self.assertLess(metrics.loss, prev_loss)
prev_loss = metrics.loss
def test_get_model_weights(self):
iterative_process = federated_averaging.build_federated_averaging_process(
model_fn=model_examples.LinearRegression,
client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=
0.1))
num_clients = 3
ds = tf.data.Dataset.from_tensor_slices(
collections.OrderedDict(
x=[[1.0, 2.0], [3.0, 4.0]],
y=[[5.0], [6.0]],
)).batch(2)
datasets = [ds] * num_clients
state = iterative_process.initialize()
self.assertIsInstance(iterative_process.get_model_weights(state),
model_utils.ModelWeights)
self.assertAllClose(
state.model.trainable,
iterative_process.get_model_weights(state).trainable)
for _ in range(3):
state, _ = iterative_process.next(state, datasets)
self.assertIsInstance(iterative_process.get_model_weights(state),
model_utils.ModelWeights)
self.assertAllClose(
state.model.trainable,
iterative_process.get_model_weights(state).trainable)
def benchmark_simple_execution(self, executor_id):
num_clients = 10
num_client_samples = 20
batch_size = 4
num_rounds = 10
ds = tf.data.Dataset.from_tensor_slices(
collections.OrderedDict(x=[[1., 2.]] * num_client_samples,
y=[[5.]] *
num_client_samples)).batch(batch_size)
federated_ds = [ds] * num_clients
building_time_array = []
build_time_start = time.time()
iterative_process = federated_averaging.build_federated_averaging_process(
model_fn=model_examples.TrainableLinearRegression)
build_time_stop = time.time()
building_time_array.append(build_time_stop - build_time_start)
name = ('computation_building_time, simple execution '
'TrainableLinearRegression, executor {}'.format(executor_id))
self.report_benchmark(name=name,
wall_time=np.mean(building_time_array),
iters=1)
initialization_array = []
initialization_start = time.time()
initial_state = iterative_process.initialize()
initialization_stop = time.time()
initialization_array.append(initialization_stop - initialization_start)
name = ('computation_initialization_time, simple execution '
'TrainableLinearRegression, executor {}'.format(executor_id))
self.report_benchmark(name=name,
wall_time=np.mean(initialization_array),
iters=1)
next_state = initial_state
execution_array = []
for _ in range(num_rounds - 1):
round_start = time.time()
next_state, _ = iterative_process.next(next_state, federated_ds)
round_stop = time.time()
execution_array.append(round_stop - round_start)
name = (
'Average per round time, {clients} clients, {examples} examples '
'per client, batch size {batch_size}, TrainableLinearRegression, '
'executor {executor}'.format(clients=num_clients,
examples=num_client_samples,
batch_size=batch_size,
executor=executor_id))
self.report_benchmark(name=name,
wall_time=np.mean(execution_array),
iters=num_rounds,
extras={'std_dev': np.std(execution_array)})
def benchmark_simple_execution(self):
num_clients = 10
num_client_samples = 20
batch_size = 4
num_rounds = 10
ds = tf.data.Dataset.from_tensor_slices({
"x": [[1., 2.]] * num_client_samples,
"y": [[5.]] * num_client_samples
}).batch(batch_size)
federated_ds = [ds] * num_clients
building_time_array = []
build_time_start = time.time()
iterative_process = federated_averaging.build_federated_averaging_process(
model_fn=model_examples.TrainableLinearRegression)
build_time_stop = time.time()
building_time_array.append(build_time_stop - build_time_start)
self.report_benchmark(
name="computation_building_time, simple execution "
"TrainableLinearRegression",
wall_time=np.mean(building_time_array),
iters=1)
initialization_array = []
initialization_start = time.time()
server_state = iterative_process.initialize()
initialization_stop = time.time()
initialization_array.append(initialization_stop - initialization_start)
self.report_benchmark(
name="computation_initialization_time, simple execution "
"TrainableLinearRegression",
wall_time=np.mean(initialization_array),
iters=1)
next_state = server_state
execution_array = []
next_state, _ = iterative_process.next(server_state, federated_ds)
for _ in range(num_rounds - 1):
round_start = time.time()
next_state, _ = iterative_process.next(next_state, federated_ds)
round_stop = time.time()
execution_array.append(round_stop - round_start)
self.report_benchmark(name="Time to execute {} rounds, {} clients, "
"{} examples per client, batch size {}, "
"TrainableLinearRegression".format(
num_rounds, num_clients, num_client_samples,
batch_size),
wall_time=np.mean(execution_array),
iters=num_rounds,
extras={"std_dev": np.std(execution_array)})
def test_orchestration_typecheck(self):
iterative_process = federated_averaging.build_federated_averaging_process(
model_fn=model_examples.TrainableLinearRegression)
expected_model_weights_type = model_utils.ModelWeights(
collections.OrderedDict([('a', tff.TensorType(tf.float32, [2, 1])),
('b', tf.float32)]),
collections.OrderedDict([('c', tf.float32)]))
# ServerState consists of a model and optimizer_state. The optimizer_state
# is provided by TensorFlow, TFF doesn't care what the actual value is.
expected_federated_server_state_type = tff.FederatedType(
optimizer_utils.ServerState(expected_model_weights_type,
test.AnyType()),
placement=tff.SERVER,
all_equal=True)
expected_federated_dataset_type = tff.FederatedType(
tff.SequenceType(
model_examples.TrainableLinearRegression.make_batch(
tff.TensorType(tf.float32, [None, 2]),
tff.TensorType(tf.float32, [None, 1]))),
tff.CLIENTS,
all_equal=False)
expected_model_output_types = tff.FederatedType(
collections.OrderedDict([
('loss', tff.TensorType(tf.float32, [])),
('num_examples', tff.TensorType(tf.int32, [])),
]),
tff.SERVER,
all_equal=True)
# `initialize` is expected to be a funcion of no arguments to a ServerState.
self.assertEqual(
tff.FunctionType(
parameter=None, result=expected_federated_server_state_type),
iterative_process.initialize.type_signature)
# `next` is expected be a function of (ServerState, Datasets) to
# ServerState.
self.assertEqual(
tff.FunctionType(
parameter=[
expected_federated_server_state_type,
expected_federated_dataset_type
],
result=(expected_federated_server_state_type,
expected_model_output_types)),
iterative_process.next.type_signature)
def benchmark_learning_keras_model_mnist(self, executor_id):
"""Code adapted from MNIST learning tutorial ipynb."""
federated_train_data = generate_fake_mnist_data()
n_rounds = 10
computation_building_start = time.time()
# pylint: disable=missing-docstring
def model_fn():
model = tf.keras.models.Sequential([
tf.keras.layers.Flatten(input_shape=(784,)),
tf.keras.layers.Dense(
10,
kernel_initializer="zeros",
bias_initializer="zeros",
activation=tf.nn.softmax)
])
model.compile(
loss=tf.keras.losses.SparseCategoricalCrossentropy(),
optimizer=tf.keras.optimizers.SGD(0.1))
return keras_utils.from_compiled_keras_model(model,
federated_train_data[0][0])
iterative_process = federated_averaging.build_federated_averaging_process(
model_fn)
computation_building_stop = time.time()
building_time = computation_building_stop - computation_building_start
self.report_benchmark(
name="computation_building_time, "
"tff.learning Keras model, executor {}".format(executor_id),
wall_time=building_time,
iters=1)
state = iterative_process.initialize()
execution_array = []
for _ in range(n_rounds):
execution_start = time.time()
state, _ = iterative_process.next(state, federated_train_data)
execution_stop = time.time()
execution_array.append(execution_stop - execution_start)
self.report_benchmark(
name="Average per round execution time, "
"tff.learning Keras model, executor {}".format(executor_id),
wall_time=np.mean(execution_array),
iters=n_rounds,
extras={"std_dev": np.std(execution_array)})
def benchmark_learning_keras_model_mnist(self, executor_id):
"""Code adapted from MNIST learning tutorial ipynb."""
federated_train_data = generate_fake_mnist_data()
x_type = tf.TensorSpec(shape=(None, 784), dtype=tf.float32)
y_type = tf.TensorSpec(shape=(None,), dtype=tf.int32)
input_spec = collections.OrderedDict(x=x_type, y=y_type)
n_rounds = 10
computation_building_start = time.time()
# pylint: disable=missing-docstring
def model_fn():
model = tf.keras.models.Sequential([
tf.keras.layers.Flatten(input_shape=(784,)),
tf.keras.layers.Dense(
10,
kernel_initializer='zeros',
bias_initializer='zeros',
activation=tf.nn.softmax)
])
return keras_utils.from_keras_model(
model,
input_spec=input_spec,
loss=tf.keras.losses.SparseCategoricalCrossentropy())
iterative_process = federated_averaging.build_federated_averaging_process(
model_fn, client_optimizer_fn=lambda: tf.keras.optimizers.SGD(0.1))
computation_building_stop = time.time()
building_time = computation_building_stop - computation_building_start
name = ('computation_building_time, tff.learning Keras model, executor {}'
.format(executor_id))
self.report_benchmark(name=name, wall_time=building_time, iters=1)
state = iterative_process.initialize()
execution_array = []
for _ in range(n_rounds):
execution_start = time.time()
state, _ = iterative_process.next(state, federated_train_data)
execution_stop = time.time()
execution_array.append(execution_stop - execution_start)
name = ('Average per round execution time, tff.learning Keras model, '
'executor {}'.format(executor_id))
self.report_benchmark(
name=name,
wall_time=np.mean(execution_array),
iters=n_rounds,
extras={'std_dev': np.std(execution_array)})
def test_basic_orchestration_execute(self):
iterative_process = federated_averaging.build_federated_averaging_process(
model_fn=model_examples.LinearRegression,
client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=
0.1))
ds = tf.data.Dataset.from_tensor_slices(
collections.OrderedDict(
x=[[1.0, 2.0], [3.0, 4.0]],
y=[[5.0], [6.0]],
)).batch(2)
num_clients = 3
self._run_test(iterative_process,
datasets=[ds] * num_clients,
expected_num_examples=2 * num_clients)
def test_recommended_aggregations_execute(self, default_aggregation):
process = federated_averaging.build_federated_averaging_process(
model_fn=model_examples.LinearRegression,
client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=
0.1),
model_update_aggregation_factory=default_aggregation())
ds = tf.data.Dataset.from_tensor_slices(
collections.OrderedDict(
x=[[1.0, 2.0], [3.0, 4.0]],
y=[[5.0], [6.0]],
)).batch(2)
num_clients = 3
state = process.initialize()
state, metrics = process.next(state, [ds] * num_clients)
self.assertNotEmpty(metrics['aggregation'])
def test_orchestration_execute(self):
iterative_process = federated_averaging.build_federated_averaging_process(
model_fn=model_examples.TrainableLinearRegression)
ds = tf.data.Dataset.from_tensor_slices({
'x': [[1., 2.], [3., 4.]],
'y': [[5.], [6.]]
}).batch(2)
federated_ds = [ds] * 3
server_state = iterative_process.initialize()
prev_loss = np.inf
for _ in range(3):
server_state, metric_outputs = iterative_process.next(
server_state, federated_ds)
self.assertEqual(metric_outputs.num_examples, 2 * len(federated_ds))
self.assertLess(metric_outputs.loss, prev_loss)
prev_loss = metric_outputs.loss
def test_execute_empty_data(self, client_optimizer):
iterative_process = federated_averaging.build_federated_averaging_process(
model_fn=model_examples.LinearRegression,
client_optimizer_fn=client_optimizer())
# Results in empty dataset with correct types and shapes.
ds = tf.data.Dataset.from_tensor_slices(
collections.OrderedDict(
x=[[1.0, 2.0]],
y=[[5.0]],
)).batch(5, drop_remainder=True)
server_state = iterative_process.initialize()
first_state, metric_outputs = iterative_process.next(
server_state, [ds] * 2)
self.assertAllClose(list(first_state.model.trainable),
[[[0.0], [0.0]], 0.0])
self.assertEqual(metric_outputs['train']['num_examples'], 0)
self.assertTrue(tf.math.is_nan(metric_outputs['train']['loss']))
def test_execute_empty_data(self):
iterative_process = federated_averaging.build_federated_averaging_process(
model_fn=model_examples.TrainableLinearRegression)
# Results in empty dataset with correct types and shapes.
ds = tf.data.Dataset.from_tensor_slices(
collections.OrderedDict([
('x', [[1.0, 2.0]]),
('y', [[5.0]]),
])).batch(5, drop_remainder=True)
federated_ds = [ds] * 2
server_state = iterative_process.initialize()
first_state, metric_outputs = iterative_process.next(
server_state, federated_ds)
self.assertAllClose(list(first_state.model.trainable),
[[[0.0], [0.0]], 0.0])
self.assertEqual(metric_outputs.num_examples, 0)
self.assertTrue(tf.math.is_nan(metric_outputs.loss))
def test_orchestration_execute(self):
iterative_process = federated_averaging.build_federated_averaging_process(
model_fn=model_examples.LinearRegression,
client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.1))
ds = tf.data.Dataset.from_tensor_slices(
collections.OrderedDict(
x=[[1.0, 2.0], [3.0, 4.0]],
y=[[5.0], [6.0]],
)).batch(2)
federated_ds = [ds] * 3
server_state = iterative_process.initialize()
prev_loss = np.inf
for _ in range(3):
server_state, metric_outputs = iterative_process.next(
server_state, federated_ds)
self.assertEqual(metric_outputs.num_examples, 2 * len(federated_ds))
self.assertLess(metric_outputs.loss, prev_loss)
prev_loss = metric_outputs.loss
def test_orchestration_execute_from_keras_with_lookup(
self, client_optimizer):
ds = tf.data.Dataset.from_tensor_slices(
collections.OrderedDict(x=[['R'], ['G'], ['B']],
y=[[1.0], [2.0], [3.0]])).batch(2)
def model_fn():
keras_model = model_examples.build_lookup_table_keras_model()
return keras_utils.from_keras_model(
keras_model,
loss=tf.keras.losses.MeanSquaredError(),
input_spec=ds.element_spec,
metrics=[NumExamplesCounter()])
iterative_process = federated_averaging.build_federated_averaging_process(
model_fn=model_fn, client_optimizer_fn=client_optimizer())
num_clients = 3
self._run_test(iterative_process,
datasets=[ds] * num_clients,
expected_num_examples=3 * num_clients)
def test_execute_empty_data(self):
iterative_process = federated_averaging.build_federated_averaging_process(
model_fn=model_examples.TrainableLinearRegression)
# Results in empty dataset with correct types and shapes.
ds = tf.data.Dataset.from_tensor_slices({
'x': [[1., 2.]],
'y': [[5.]]
}).batch(5, drop_remainder=True)
federated_ds = [ds] * 2
server_state = iterative_process.initialize()
first_state, metric_outputs = iterative_process.next(
server_state, federated_ds)
self.assertEqual(
self.evaluate(tf.reduce_sum(first_state.model.trainable.a)) +
self.evaluate(tf.reduce_sum(first_state.model.trainable.b)), 0)
self.assertEqual(metric_outputs.num_examples, 0)
self.assertTrue(tf.is_nan(metric_outputs.loss))
def test_orchestration_execute_from_keras(self, build_keras_model_fn):
ds = tf.data.Dataset.from_tensor_slices(
collections.OrderedDict(
x=[[1.0, 2.0], [3.0, 4.0]],
y=[[5.0], [6.0]],
)).batch(2)
def model_fn():
keras_model = build_keras_model_fn(feature_dims=2)
return keras_utils.from_keras_model(
keras_model,
loss=tf.keras.losses.MeanSquaredError(),
input_spec=ds.element_spec,
metrics=[NumExamplesCounter()])
iterative_process = federated_averaging.build_federated_averaging_process(
model_fn=model_fn,
client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=
0.01))
num_clients = 3
self._run_test(iterative_process,
datasets=[ds] * num_clients,
expected_num_examples=2 * num_clients)
