def test_self_contained_example(self):
client_data = create_client_data()
model = MnistModel()
optimizer_fn = lambda: tf.keras.optimizers.SGD(learning_rate=0.1)
losses = []
for r in range(2):
optimizer = optimizer_fn()
simple_fedavg_tff._initialize_optimizer_vars(model, optimizer)
server_message = simple_fedavg_tf.BroadcastMessage(
model_weights=model.weights, round_num=r)
outputs = simple_fedavg_tf.client_update(model, client_data(),
server_message, optimizer)
losses.append(outputs.model_output.numpy())
self.assertAllEqual(int(outputs.client_weight.numpy()), 2)
self.assertLess(losses[1], losses[0])
def client_update_fn(tf_dataset, server_message):
model = model_fn()
client_optimizer = client_optimizer_fn()
return client_update(model, tf_dataset, server_message,
client_optimizer)