<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Copyright 2020 The TensorFlow Authors.
Step1:  <table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https
Step2:  NOTE
Step4:  In order to feed the dataset into our model, we flatten the data, and convert each example into a tuple of the form (flattened_image_vector, label).
Step5:  We now select a small number of clients, and apply the preprocessing above to their datasets.
Step6:  Preparing the model
We use the same model as in the image classification tutorial. This model (implemented via tf.keras) has a single hidden layer, followed by a softmax layer.
Step7:  In order to use this model in TFF, we wrap the Keras model as a tff.learning.Model. This allows us to perform the model's forward pass within TFF, and extract model outputs. For more details, also see the image classification tutorial.
Step8:  While we used tf.keras to create a tff.learning.Model, TFF supports much more general models. These models have the following relevant attributes capturing the model weights
Step9:  This function looks good, but as we will see later, we will need to make a small modification to make it a "TFF computation".
We also want to sketch the next_fn.
Step11:  We'll focus on implementing these four components separately. We first focus on the parts that can be implemented in pure TensorFlow, namely the client and server update steps.
TensorFlow Blocks
Client update
We will use our tff.learning.Model to do client training in essentially the same way you would train a TensorFlow model. In particular, we will use tf.GradientTape to compute the gradient on batches of data, then apply these gradient using a client_optimizer. We focus only on the trainable weights.
Step13:  Server Update
The server update for FedAvg is simpler than the client update. We will implement "vanilla" federated averaging, in which we simply replace the server model weights by the average of the client model weights. Again, we only focus on the trainable weights.
Step14:  The snippet could be simplified by simply returning the mean_client_weights. However, more advanced implementations of Federated Averaging use mean_client_weights with more sophisticated techniques, such as momentum or adaptivity.
Challenge
Step15:  Federated types are specified by a type T of its member constituents (eg. tf.float32) and a group G of devices. We will focus on the cases where G is either tff.CLIENTS or tff.SERVER. Such a federated type is represented as {T}@G, as shown below.
Step16:  Why do we care so much about placements? A key goal of TFF is to enable writing code that could be deployed on a real distributed system. This means that it is vital to reason about which subsets of devices execute which code, and where different pieces of data reside.
TFF focuses on three things
Step17:  You might ask, how is this different from the tf.function decorator in TensorFlow? The key answer is that the code generated by tff.federated_computation is neither TensorFlow nor Python code; It is a specification of a distributed system in an internal platform-independent glue language.
While this may sound complicated, you can think of TFF computations as functions with well-defined type signatures. These type signatures can be directly queried.
Step18:  This tff.federated_computation accepts arguments of federated type {float32}@CLIENTS, and returns values of federated type {float32}@SERVER. Federated computations may also go from server to client, from client to client, or from server to server. Federated computations can also be composed like normal functions, as long as their type signatures match up.
To support development, TFF allows you to invoke a tff.federated_computation as a Python function. For example, we can call
Step19:  Non-eager computations and TensorFlow
There are two key restrictions to be aware of. First, when the Python interpreter encounters a tff.federated_computation decorator, the function is traced once and serialized for future use. Due to the decentralized nature of Federated Learning, this future usage may occur elsewhere, such as a remote execution environment. Therefore, TFF computations are fundamentally non-eager. This behavior is somewhat analogous to that of the tf.function decorator in TensorFlow.
Second, a federated computation can only consist of federated operators (such as tff.federated_mean), they cannot contain TensorFlow operations. TensorFlow code must be confined to blocks decorated with tff.tf_computation. Most ordinary TensorFlow code can be directly decorated, such as the following function that takes a number and adds 0.5 to it.
Step20:  These also have type signatures, but without placements. For example, we can call
Step21:  Here we see an important difference between tff.federated_computation and tff.tf_computation. The former has explicit placements, while the latter does not.
We can use tff.tf_computation blocks in federated computations by specifying placements. Let's create a function that adds half, but only to federated floats at the clients. We can do this by using tff.federated_map, which applies a given tff.tf_computation, while preserving the placement.
Step22:  This function is almost identical to add_half, except that it only accepts values with placement at tff.CLIENTS, and returns values with the same placement. We can see this in its type signature
Step23:  In summary
Step24:  We can then pass this directly into a federated computation using tff.federated_value.
Step25:  Creating the next_fn
We now use our client and server update code to write the actual algorithm. We will first turn our client_update into a tff.tf_computation that accepts a client datasets and server weights, and outputs an updated client weights tensor.
We will need the corresponding types to properly decorate our function. Luckily, the type of the server weights can be extracted directly from our model.
Step26:  Let's look at the dataset type signature. Remember that we took 28 by 28 images (with integer labels) and flattened them.
Step27:  We can also extract the model weights type by using our server_init function above.
Step28:  Examining the type signature, we'll be able to see the architecture of our model!
Step29:  We can now create our tff.tf_computation for the client update.
Step30:  The tff.tf_computation version of the server update can be defined in a similar way, using types we've already extracted.
Step31:  Last, but not least, we need to create the tff.federated_computation that brings this all together. This function will accept two federated values, one corresponding to the server weights (with placement tff.SERVER), and the other corresponding to the client datasets (with placement tff.CLIENTS).
Note that both these types were defined above! We simply need to give them the proper placement using tff.FederatedType.
Step32:  Remember the 4 elements of an FL algorithm?
A server-to-client broadcast step.
A local client update step.
A client-to-server upload step.
A server update step.
Now that we've built up the above, each part can be compactly represented as a single line of TFF code. This simplicity is why we had to take extra care to specify things such as federated types!
Step33:  We now have a tff.federated_computation for both the algorithm initialization, and for running one step of the algorithm. To finish our algorithm, we pass these into tff.templates.IterativeProcess.
Step34:  Let's look at the type signature of the initialize and next functions of our iterative process.
Step35:  This reflects the fact that federated_algorithm.initialize is a no-arg function that returns a single-layer model (with a 784-by-10 weight matrix, and 10 bias units).
Step36:  Here, we see that federated_algorithm.next accepts a server model and client data, and returns an updated server model.
Evaluating the algorithm
Let's run a few rounds, and see how the loss changes. First, we will define an evaluation function using the centralized approach discussed in the second tutorial.
We first create a centralized evaluation dataset, and then apply the same preprocessing we used for the training data.
Step37:  Next, we write a function that accepts a server state, and uses Keras to evaluate on the test dataset. If you're familiar with tf.Keras, this will all look familiar, though note the use of set_weights!
Step38:  Now, let's initialize our algorithm and evaluate on the test set.
Step39:  Let's train for a few rounds and see if anything changes.
<ASSISTANT_TASK:>
Python Code:
#@title 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
#
# https://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.
Explanation: Copyright 2020 The TensorFlow Authors.
End of explanation
#@test {"skip": true}
!pip install --quiet --upgrade tensorflow-federated
!pip install --quiet --upgrade nest-asyncio
import nest_asyncio
nest_asyncio.apply()
import tensorflow as tf
import tensorflow_federated as tff
Explanation: <table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https://www.tensorflow.org/federated/tutorials/building_your_own_federated_learning_algorithm"><img src="https://www.tensorflow.org/images/tf_logo_32px.png" />View on TensorFlow.org</a>
  </td>
  <td>
    <a target="_blank" href="https://colab.research.google.com/github/tensorflow/federated/blob/v0.27.0/docs/tutorials/building_your_own_federated_learning_algorithm.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png" />Run in Google Colab</a>
  </td>
  <td>
    <a target="_blank" href="https://github.com/tensorflow/federated/blob/v0.27.0/docs/tutorials/building_your_own_federated_learning_algorithm.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View source on GitHub</a>
  </td>
  <td>
    <a href="https://storage.googleapis.com/tensorflow_docs/federated/docs/tutorials/building_your_own_federated_learning_algorithm.ipynb"><img src="https://www.tensorflow.org/images/download_logo_32px.png" />Download notebook</a>
  </td>
</table>
Before we start
Before we start, please run the following to make sure that your environment is
correctly setup. If you don't see a greeting, please refer to the
Installation guide for instructions.
End of explanation
emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data()
Explanation: NOTE: This colab has been verified to work with the latest released version of the tensorflow_federated pip package, but the Tensorflow Federated project is still in pre-release development and may not work on main.
Building Your Own Federated Learning Algorithm
In the image classification and
text generation tutorials, we learned how to set up model and data pipelines for Federated Learning (FL), and performed federated training via the tff.learning API layer of TFF.
This is only the tip of the iceberg when it comes to FL research. In this tutorial, we discuss how to implement federated learning algorithms without deferring to the tff.learning API. We aim to accomplish the following:
Goals:
Understand the general structure of federated learning algorithms.
Explore the Federated Core of TFF.
Use the Federated Core to implement Federated Averaging directly.
While this tutorial is self-contained, we recommend first reading the image classification and
text generation tutorials.
Preparing the input data
We first load and preprocess the EMNIST dataset included in TFF. For more details, see the image classification tutorial.
End of explanation
NUM_CLIENTS = 10
BATCH_SIZE = 20
def preprocess(dataset):
  def batch_format_fn(element):
    Flatten a batch of EMNIST data and return a (features, label) tuple.
    return (tf.reshape(element['pixels'], [-1, 784]), 
            tf.reshape(element['label'], [-1, 1]))
  return dataset.batch(BATCH_SIZE).map(batch_format_fn)
Explanation: In order to feed the dataset into our model, we flatten the data, and convert each example into a tuple of the form (flattened_image_vector, label).
End of explanation
client_ids = sorted(emnist_train.client_ids)[:NUM_CLIENTS]
federated_train_data = [preprocess(emnist_train.create_tf_dataset_for_client(x))
  for x in client_ids
]
Explanation: We now select a small number of clients, and apply the preprocessing above to their datasets.
End of explanation
def create_keras_model():
  initializer = tf.keras.initializers.GlorotNormal(seed=0)
  return tf.keras.models.Sequential([
      tf.keras.layers.Input(shape=(784,)),
      tf.keras.layers.Dense(10, kernel_initializer=initializer),
      tf.keras.layers.Softmax(),
  ])
Explanation: Preparing the model
We use the same model as in the image classification tutorial. This model (implemented via tf.keras) has a single hidden layer, followed by a softmax layer.
End of explanation
def model_fn():
  keras_model = create_keras_model()
  return tff.learning.from_keras_model(
      keras_model,
      input_spec=federated_train_data[0].element_spec,
      loss=tf.keras.losses.SparseCategoricalCrossentropy(),
      metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])
Explanation: In order to use this model in TFF, we wrap the Keras model as a tff.learning.Model. This allows us to perform the model's forward pass within TFF, and extract model outputs. For more details, also see the image classification tutorial.
End of explanation
def initialize_fn():
  model = model_fn()
  return model.trainable_variables
Explanation: While we used tf.keras to create a tff.learning.Model, TFF supports much more general models. These models have the following relevant attributes capturing the model weights:
trainable_variables: An iterable of the tensors corresponding to trainable layers.
non_trainable_variables: An iterable of the tensors corresponding to non-trainable layers.
For our purposes, we will only use the trainable_variables. (as our model only has those!).
Building your own Federated Learning algorithm
While the tff.learning API allows one to create many variants of Federated Averaging, there are other federated algorithms that do not fit neatly into this framework. For example, you may want to add regularization, clipping, or more complicated algorithms such as federated GAN training. You may also be instead be interested in federated analytics.
For these more advanced algorithms, we'll have to write our own custom algorithm using TFF. In many cases, federated algorithms have 4 main components:
A server-to-client broadcast step.
A local client update step.
A client-to-server upload step.
A server update step.
In TFF, we generally represent federated algorithms as a tff.templates.IterativeProcess (which we refer to as just an IterativeProcess throughout). This is a class that contains initialize and next functions. Here, initialize is used to initialize the server, and next will perform one communication round of the federated algorithm. Let's write a skeleton of what our iterative process for FedAvg should look like.
First, we have an initialize function that simply creates a tff.learning.Model, and returns its trainable weights.
End of explanation
def next_fn(server_weights, federated_dataset):
  # Broadcast the server weights to the clients.
  server_weights_at_client = broadcast(server_weights)
  # Each client computes their updated weights.
  client_weights = client_update(federated_dataset, server_weights_at_client)
  # The server averages these updates.
  mean_client_weights = mean(client_weights)
  # The server updates its model.
  server_weights = server_update(mean_client_weights)
  return server_weights
Explanation: This function looks good, but as we will see later, we will need to make a small modification to make it a "TFF computation".
We also want to sketch the next_fn.
End of explanation
@tf.function
def client_update(model, dataset, server_weights, client_optimizer):
  Performs training (using the server model weights) on the client's dataset.
  # Initialize the client model with the current server weights.
  client_weights = model.trainable_variables
  # Assign the server weights to the client model.
  tf.nest.map_structure(lambda x, y: x.assign(y),
                        client_weights, server_weights)
  # Use the client_optimizer to update the local model.
  for batch in dataset:
    with tf.GradientTape() as tape:
      # Compute a forward pass on the batch of data
      outputs = model.forward_pass(batch)
    # Compute the corresponding gradient
    grads = tape.gradient(outputs.loss, client_weights)
    grads_and_vars = zip(grads, client_weights)
    # Apply the gradient using a client optimizer.
    client_optimizer.apply_gradients(grads_and_vars)
  return client_weights
Explanation: We'll focus on implementing these four components separately. We first focus on the parts that can be implemented in pure TensorFlow, namely the client and server update steps.
TensorFlow Blocks
Client update
We will use our tff.learning.Model to do client training in essentially the same way you would train a TensorFlow model. In particular, we will use tf.GradientTape to compute the gradient on batches of data, then apply these gradient using a client_optimizer. We focus only on the trainable weights.
End of explanation
@tf.function
def server_update(model, mean_client_weights):
  Updates the server model weights as the average of the client model weights.
  model_weights = model.trainable_variables
  # Assign the mean client weights to the server model.
  tf.nest.map_structure(lambda x, y: x.assign(y),
                        model_weights, mean_client_weights)
  return model_weights
Explanation: Server Update
The server update for FedAvg is simpler than the client update. We will implement "vanilla" federated averaging, in which we simply replace the server model weights by the average of the client model weights. Again, we only focus on the trainable weights.
End of explanation
federated_float_on_clients = tff.FederatedType(tf.float32, tff.CLIENTS)
Explanation: The snippet could be simplified by simply returning the mean_client_weights. However, more advanced implementations of Federated Averaging use mean_client_weights with more sophisticated techniques, such as momentum or adaptivity.
Challenge: Implement a version of server_update that updates the server weights to be the midpoint of model_weights and mean_client_weights. (Note: This kind of "midpoint" approach is analogous to recent work on the Lookahead optimizer!).
So far, we've only written pure TensorFlow code. This is by design, as TFF allows you to use much of the TensorFlow code you're already familiar with. However, now we have to specify the orchestration logic, that is, the logic that dictates what the server broadcasts to the client, and what the client uploads to the server.
This will require the Federated Core of TFF.
Introduction to the Federated Core
The Federated Core (FC) is a set of lower-level interfaces that serve as the foundation for the tff.learning API. However, these interfaces are not limited to learning. In fact, they can be used for analytics and many other computations over distributed data.
At a high-level, the federated core is a development environment that enables compactly expressed program logic to combine TensorFlow code with distributed communication operators (such as distributed sums and broadcasts). The goal is to give researchers and practitioners expliict control over the distributed communication in their systems, without requiring system implementation details (such as specifying point-to-point network message exchanges).
One key point is that TFF is designed for privacy-preservation. Therefore, it allows explicit control over where data resides, to prevent unwanted accumulation of data at the centralized server location.
Federated data
A key concept in TFF is "federated data", which refers to a collection of data items hosted across a group of devices in a distributed system (eg. client datasets, or the server model weights). We model the entire collection of data items across all devices as a single federated value.
For example, suppose we have client devices that each have a float representing the temperature of a sensor. We could represent it as a federated float by
End of explanation
str(federated_float_on_clients)
Explanation: Federated types are specified by a type T of its member constituents (eg. tf.float32) and a group G of devices. We will focus on the cases where G is either tff.CLIENTS or tff.SERVER. Such a federated type is represented as {T}@G, as shown below.
End of explanation
@tff.federated_computation(tff.FederatedType(tf.float32, tff.CLIENTS))
def get_average_temperature(client_temperatures):
  return tff.federated_mean(client_temperatures)
Explanation: Why do we care so much about placements? A key goal of TFF is to enable writing code that could be deployed on a real distributed system. This means that it is vital to reason about which subsets of devices execute which code, and where different pieces of data reside.
TFF focuses on three things: data, where the data is placed, and how the data is being transformed. The first two are encapsulated in federated types, while the last is encapsulated in federated computations.
Federated computations
TFF is a strongly-typed functional programming environment whose basic units are federated computations. These are pieces of logic that accept federated values as input, and return federated values as output.
For example, suppose we wanted to average the temperatures on our client sensors. We could define the following (using our federated float):
End of explanation
str(get_average_temperature.type_signature)
Explanation: You might ask, how is this different from the tf.function decorator in TensorFlow? The key answer is that the code generated by tff.federated_computation is neither TensorFlow nor Python code; It is a specification of a distributed system in an internal platform-independent glue language.
While this may sound complicated, you can think of TFF computations as functions with well-defined type signatures. These type signatures can be directly queried.
End of explanation
get_average_temperature([68.5, 70.3, 69.8])
Explanation: This tff.federated_computation accepts arguments of federated type {float32}@CLIENTS, and returns values of federated type {float32}@SERVER. Federated computations may also go from server to client, from client to client, or from server to server. Federated computations can also be composed like normal functions, as long as their type signatures match up.
To support development, TFF allows you to invoke a tff.federated_computation as a Python function. For example, we can call
End of explanation
@tff.tf_computation(tf.float32)
def add_half(x):
  return tf.add(x, 0.5)
Explanation: Non-eager computations and TensorFlow
There are two key restrictions to be aware of. First, when the Python interpreter encounters a tff.federated_computation decorator, the function is traced once and serialized for future use. Due to the decentralized nature of Federated Learning, this future usage may occur elsewhere, such as a remote execution environment. Therefore, TFF computations are fundamentally non-eager. This behavior is somewhat analogous to that of the tf.function decorator in TensorFlow.
Second, a federated computation can only consist of federated operators (such as tff.federated_mean), they cannot contain TensorFlow operations. TensorFlow code must be confined to blocks decorated with tff.tf_computation. Most ordinary TensorFlow code can be directly decorated, such as the following function that takes a number and adds 0.5 to it.
End of explanation
str(add_half.type_signature)
Explanation: These also have type signatures, but without placements. For example, we can call
End of explanation
@tff.federated_computation(tff.FederatedType(tf.float32, tff.CLIENTS))
def add_half_on_clients(x):
  return tff.federated_map(add_half, x)
Explanation: Here we see an important difference between tff.federated_computation and tff.tf_computation. The former has explicit placements, while the latter does not.
We can use tff.tf_computation blocks in federated computations by specifying placements. Let's create a function that adds half, but only to federated floats at the clients. We can do this by using tff.federated_map, which applies a given tff.tf_computation, while preserving the placement.
End of explanation
str(add_half_on_clients.type_signature)
Explanation: This function is almost identical to add_half, except that it only accepts values with placement at tff.CLIENTS, and returns values with the same placement. We can see this in its type signature:
End of explanation
@tff.tf_computation
def server_init():
  model = model_fn()
  return model.trainable_variables
Explanation: In summary:
TFF operates on federated values.
Each federated value has a federated type, with a type (eg. tf.float32) and a placement (eg. tff.CLIENTS).
Federated values can be transformed using federated computations, which must be decorated with tff.federated_computation and a federated type signature.
TensorFlow code must be contained in blocks with tff.tf_computation decorators. 
These blocks can then be incorporated into federated computations.
Building your own Federated Learning algorithm, revisited
Now that we've gotten a glimpse of the Federated Core, we can build our own federated learning algorithm. Remember that above, we defined an initialize_fn and next_fn for our algorithm. The next_fn will make use of the client_update and server_update we defined using pure TensorFlow code.
However, in order to make our algorithm a federated computation, we will need both the next_fn and initialize_fn to each be a tff.federated_computation.
TensorFlow Federated blocks
Creating the initialization computation
The initialize function will be quite simple: We will create a model using model_fn. However, remember that we must separate out our TensorFlow code using tff.tf_computation.
End of explanation
@tff.federated_computation
def initialize_fn():
  return tff.federated_value(server_init(), tff.SERVER)
Explanation: We can then pass this directly into a federated computation using tff.federated_value.
End of explanation
whimsy_model = model_fn()
tf_dataset_type = tff.SequenceType(whimsy_model.input_spec)
Explanation: Creating the next_fn
We now use our client and server update code to write the actual algorithm. We will first turn our client_update into a tff.tf_computation that accepts a client datasets and server weights, and outputs an updated client weights tensor.
We will need the corresponding types to properly decorate our function. Luckily, the type of the server weights can be extracted directly from our model.
End of explanation
str(tf_dataset_type)
Explanation: Let's look at the dataset type signature. Remember that we took 28 by 28 images (with integer labels) and flattened them.
End of explanation
model_weights_type = server_init.type_signature.result
Explanation: We can also extract the model weights type by using our server_init function above.
End of explanation
str(model_weights_type)
Explanation: Examining the type signature, we'll be able to see the architecture of our model!
End of explanation
@tff.tf_computation(tf_dataset_type, model_weights_type)
def client_update_fn(tf_dataset, server_weights):
  model = model_fn()
  client_optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)
  return client_update(model, tf_dataset, server_weights, client_optimizer)
Explanation: We can now create our tff.tf_computation for the client update.
End of explanation
@tff.tf_computation(model_weights_type)
def server_update_fn(mean_client_weights):
  model = model_fn()
  return server_update(model, mean_client_weights)
Explanation: The tff.tf_computation version of the server update can be defined in a similar way, using types we've already extracted.
End of explanation
federated_server_type = tff.FederatedType(model_weights_type, tff.SERVER)
federated_dataset_type = tff.FederatedType(tf_dataset_type, tff.CLIENTS)
Explanation: Last, but not least, we need to create the tff.federated_computation that brings this all together. This function will accept two federated values, one corresponding to the server weights (with placement tff.SERVER), and the other corresponding to the client datasets (with placement tff.CLIENTS).
Note that both these types were defined above! We simply need to give them the proper placement using tff.FederatedType.
End of explanation
@tff.federated_computation(federated_server_type, federated_dataset_type)
def next_fn(server_weights, federated_dataset):
  # Broadcast the server weights to the clients.
  server_weights_at_client = tff.federated_broadcast(server_weights)
  # Each client computes their updated weights.
  client_weights = tff.federated_map(
      client_update_fn, (federated_dataset, server_weights_at_client))
  
  # The server averages these updates.
  mean_client_weights = tff.federated_mean(client_weights)
  # The server updates its model.
  server_weights = tff.federated_map(server_update_fn, mean_client_weights)
  return server_weights
Explanation: Remember the 4 elements of an FL algorithm?
A server-to-client broadcast step.
A local client update step.
A client-to-server upload step.
A server update step.
Now that we've built up the above, each part can be compactly represented as a single line of TFF code. This simplicity is why we had to take extra care to specify things such as federated types!
End of explanation
federated_algorithm = tff.templates.IterativeProcess(
    initialize_fn=initialize_fn,
    next_fn=next_fn
)
Explanation: We now have a tff.federated_computation for both the algorithm initialization, and for running one step of the algorithm. To finish our algorithm, we pass these into tff.templates.IterativeProcess.
End of explanation
str(federated_algorithm.initialize.type_signature)
Explanation: Let's look at the type signature of the initialize and next functions of our iterative process.
End of explanation
str(federated_algorithm.next.type_signature)
Explanation: This reflects the fact that federated_algorithm.initialize is a no-arg function that returns a single-layer model (with a 784-by-10 weight matrix, and 10 bias units).
End of explanation
central_emnist_test = emnist_test.create_tf_dataset_from_all_clients()
central_emnist_test = preprocess(central_emnist_test)
Explanation: Here, we see that federated_algorithm.next accepts a server model and client data, and returns an updated server model.
Evaluating the algorithm
Let's run a few rounds, and see how the loss changes. First, we will define an evaluation function using the centralized approach discussed in the second tutorial.
We first create a centralized evaluation dataset, and then apply the same preprocessing we used for the training data.
End of explanation
def evaluate(server_state):
  keras_model = create_keras_model()
  keras_model.compile(
      loss=tf.keras.losses.SparseCategoricalCrossentropy(),
      metrics=[tf.keras.metrics.SparseCategoricalAccuracy()]  
  )
  keras_model.set_weights(server_state)
  keras_model.evaluate(central_emnist_test)
Explanation: Next, we write a function that accepts a server state, and uses Keras to evaluate on the test dataset. If you're familiar with tf.Keras, this will all look familiar, though note the use of set_weights!
End of explanation
server_state = federated_algorithm.initialize()
evaluate(server_state)
Explanation: Now, let's initialize our algorithm and evaluate on the test set.
End of explanation
for round in range(15):
  server_state = federated_algorithm.next(server_state, federated_train_data)
evaluate(server_state)
Explanation: Let's train for a few rounds and see if anything changes.
End of explanation
<END_TASK>