def test_client_tf(self):
model = self.model()
dataset = self.dataset()
client_tf = federated_averaging.ClientFedAvg(model)
init_op = tf.group(
model_utils.model_initializer(model),
tf.compat.v1.initializers.variables(client_tf.variables),
name='fedavg_initializer')
client_outputs = client_tf(dataset, self.initial_weights())
tf.compat.v1.get_default_graph().finalize()
with self.session() as sess:
sess.run(init_op)
out = sess.run(client_outputs)
# Both trainable parameters should have been updated,
# and we don't return the non-trainable 'c'.
self.assertCountEqual(['a', 'b'], list(out.weights_delta.keys()))
self.assertGreater(np.linalg.norm(out.weights_delta['a']), 0.1)
self.assertGreater(np.linalg.norm(out.weights_delta['b']), 0.1)
self.assertEqual(out.weights_delta_weight, 8.0)
self.assertEqual(out.optimizer_output['num_examples'], 8)
self.assertEqual(out.optimizer_output['has_non_finite_delta'], 0)
self.assertEqual(out.model_output['num_examples'], 8)
self.assertEqual(out.model_output['num_batches'], 3)
self.assertBetween(out.model_output['loss'],
np.finfo(np.float32).eps, 10.0)
def test_client_tf_custom_delta_weight(self):
model = self.model()
dataset = self.dataset()
client_tf = federated_averaging.ClientFedAvg(
model, client_weight_fn=lambda _: tf.constant(1.5))
out = client_tf(dataset, self.initial_weights())
self.assertEqual(self.evaluate(out.weights_delta_weight), 1.5)
def test_client_tf(self):
model = self.create_model()
dataset = self.create_dataset()
client_tf = federated_averaging.ClientFedAvg(
model, tf.keras.optimizers.SGD(learning_rate=0.1))
client_outputs = self.evaluate(
client_tf(dataset, self.initial_weights()))
# Both trainable parameters should have been updated,
# and we don't return the non-trainable variable.
self.assertAllGreater(
np.linalg.norm(client_outputs.weights_delta, axis=-1), 0.1)
self.assertEqual(client_outputs.weights_delta_weight, 8.0)
self.assertEqual(client_outputs.optimizer_output['num_examples'], 8)
self.assertEqual(
client_outputs.optimizer_output['has_non_finite_delta'], 0)
self.assertDictContainsSubset(
{
'num_examples': 8,
'num_examples_float': 8.0,
'num_batches': 3,
}, client_outputs.model_output)
self.assertBetween(client_outputs.model_output['loss'],
np.finfo(np.float32).eps, 10.0)
def test_client_tf(self, weighted, simulation, optimizer_kwargs,
expected_norm):
model = self.create_model()
dataset = self.create_dataset()
if weighted:
client_weighting = client_weight_lib.ClientWeighting.NUM_EXAMPLES
else:
client_weighting = client_weight_lib.ClientWeighting.UNIFORM
client_tf = federated_averaging.ClientFedAvg(
model,
tf.keras.optimizers.SGD(learning_rate=0.1, **optimizer_kwargs),
client_weighting=client_weighting,
use_experimental_simulation_loop=simulation)
client_outputs = self.evaluate(
client_tf(dataset, self.initial_weights()))
# Both trainable parameters should have been updated, and we don't return
# the non-trainable variable.
self.assertAllGreater(
np.linalg.norm(client_outputs.weights_delta, axis=-1),
expected_norm)
if weighted:
self.assertEqual(client_outputs.weights_delta_weight, 8.0)
self.assertEqual(client_outputs.optimizer_output['num_examples'], 8)
self.assertDictContainsSubset(
{
'num_examples': 8,
'num_examples_float': 8.0,
'num_batches': 3,
}, client_outputs.model_output)
self.assertBetween(client_outputs.model_output['loss'],
np.finfo(np.float32).eps, 10.0)
def test_client_tf_custom_delta_weight(self):
model = self.create_model()
dataset = self.create_dataset()
client_tf = federated_averaging.ClientFedAvg(
model,
tf.keras.optimizers.SGD(learning_rate=0.1),
client_weight_fn=lambda _: tf.constant(1.5))
client_outputs = client_tf(dataset, self.initial_weights())
self.assertEqual(self.evaluate(client_outputs.weights_delta_weight), 1.5)
def test_client_tf_dataset_reduce_fn(self, simulation, mock_method):
model = self.create_model()
dataset = self.create_dataset()
client_tf = federated_averaging.ClientFedAvg(
model,
tf.keras.optimizers.SGD(learning_rate=0.1),
use_experimental_simulation_loop=simulation)
client_tf(dataset, self.initial_weights())
if simulation:
mock_method.assert_not_called()
else:
mock_method.assert_called()
def test_non_finite_aggregation(self, bad_value):
model = self.create_model()
dataset = self.create_dataset()
client_tf = federated_averaging.ClientFedAvg(
model, tf.keras.optimizers.SGD(learning_rate=0.1))
init_weights = self.initial_weights()
init_weights.trainable[1] = bad_value
client_outputs = client_tf(dataset, init_weights)
self.assertEqual(self.evaluate(client_outputs.weights_delta_weight),
0.0)
self.assertAllClose(self.evaluate(client_outputs.weights_delta),
[[[0.0], [0.0]], 0.0])
def test_non_finite_aggregation(self, bad_value):
model = self.model()
dataset = self.dataset()
client_tf = federated_averaging.ClientFedAvg(model)
init_weights = self.initial_weights()
init_weights.trainable['b'] = bad_value
out = client_tf(dataset, init_weights)
self.assertEqual(out.weights_delta_weight.numpy(), 0.0)
self.assertAllClose(out.weights_delta['a'].numpy(), np.array([[0.0],
[0.0]]))
self.assertAllClose(out.weights_delta['b'].numpy(), 0.0)
self.assertEqual(out.optimizer_output['has_non_finite_delta'].numpy(), 1)
def test_non_finite_aggregation(self, bad_value):
model = self.model()
dataset = self.dataset()
client_tf = federated_averaging.ClientFedAvg(model)
init_weights = self.initial_weights()
init_weights.trainable[1] = bad_value
client_outputs = client_tf(dataset, init_weights)
self.assertEqual(self.evaluate(client_outputs.weights_delta_weight),
0.0)
self.assertAllClose(self.evaluate(client_outputs.weights_delta),
[[[0.0], [0.0]], 0.0])
self.assertEqual(
self.evaluate(
client_outputs.optimizer_output['has_non_finite_delta']), 1)