<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 2022 The TensorFlow Authors.
Step1:  Loading Remote Data in TFF
NOTE
Step2:  Preparing the input data
Let's begin by loading TFF's federated version of the EMNIST dataset from the built-in repository
Step3:  We'll construct a preprocessing function to transform the raw examples in the EMNIST dataset from 28x28 images into 784-element arrays. Additionally, the function will shuffle the individual examples, and rename the features from pixels and label, to x and y for use with Keras. We also throw in a repeat over the data set to run several epochs.
Step4:  Let's verify this works
Step5:  We'll use the EMNIST dataset to train a model by loading and preprocessing individual clients (emulating distinct partitions) through an implementation of DataBackend.
Defining a DataBackend
We need an instance of DataBackend to instruct TFF workers (these are the processes that run the client-side of federated computations) how to load and tranform the local data stored in remote locations. A DataBackend is a programmatic construct that resolves symbolic references, represented as application-specific URIs, to materialized payloads that downstream TFF operations can process. Specifically, a DataBackend object is wrapped by a DataExecutor, which queries the object when the TFF runtime encounters an operation that fetches the data.
In this example, an Id to a client is encoded in a URI, which is parsed by our DataBackend definition to retrieve the corresponding client data, convert it to tf.Dataset, and then apply our preprocess function.
Step6:  Plugging the DataBackend into the ExecutionContext
TFF computations are invoked by an ExecutionContext and in order for data URIs defined in TFF computations to be understood at runtime, a custom context must be defined that includes a pointer to the DataBackend we just created, so URIs can be properly resolved.
The DataBackend works in tandem with DataExecutor to supply the executor with operable data that the executor can relay to requesting executors in order to complete a TFF computation.
Step7:  Training the model
Now we are ready to train a model in a federated fashion. Lets define a Keras model along with training hyperparameters
Step8:  We can pass this TFF-wrapped definition of our model
to a Federated Averaging algorithm by invoking the helper
function tff.learning.build_federated_averaging_process, as follows
Step9:  The initialize computation returns the initial state of the
Federated Averaging process.
To run a round of training, we need to construct a sample of data by organizing a sample of URI references as follows
Step10:  Now we can round a round of training
Step11:  And we can run a few more rounds
<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 2022 The TensorFlow Authors.
End of explanation
#@title Set up open-source environment
#@test {"skip": true}
!pip install --quiet --upgrade tensorflow-federated
!pip install --quiet --upgrade nest-asyncio
import nest_asyncio
nest_asyncio.apply()
#@title Import packages
import collections
import numpy as np
import tensorflow as tf
import tensorflow_federated as tff
np.random.seed(0)
Explanation: Loading Remote Data in TFF
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.
In a production environment, the raw data for a federated computation is typically distributed across machines and requires special preprocessing and loading before it's usable.
This tutorial describes how to load data stored in those remote locations with TFF's DataBackend and DataExecutor interfaces. But to keep the example simple, the dataset will exist entirely in memory and we'll smiulate the fetching as if the dataset was  partitioned over a network.
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https://www.tensorflow.org/federated/tutorials/loading_remote_data"><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/loading_remote_data.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/loading_remote_data.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/loading_remote_data.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: Preparing the input data
Let's begin by loading TFF's federated version of the EMNIST dataset from the built-in repository:
End of explanation
NUM_CLIENTS = 10
NUM_EPOCHS = 5
SHUFFLE_BUFFER = 100
def preprocess(dataset):
  def map_fn(element):
    return collections.OrderedDict(
        x=tf.reshape(element['pixels'], [-1, 784]),
        y=tf.reshape(element['label'], [-1, 1]))
  return dataset.repeat(NUM_EPOCHS).shuffle(SHUFFLE_BUFFER, seed=1).map(map_fn)
Explanation: We'll construct a preprocessing function to transform the raw examples in the EMNIST dataset from 28x28 images into 784-element arrays. Additionally, the function will shuffle the individual examples, and rename the features from pixels and label, to x and y for use with Keras. We also throw in a repeat over the data set to run several epochs.
End of explanation
example_dataset = emnist_train.create_tf_dataset_for_client(
    emnist_train.client_ids[0])
preprocessed_example_dataset = preprocess(example_dataset)
Explanation: Let's verify this works:
End of explanation
class TestDataBackend(tff.framework.DataBackend):
  async def materialize(self, data, type_spec):
    client_id = int(data.uri[-1])
    client_dataset = emnist_train.create_tf_dataset_for_client(
        emnist_train.client_ids[client_id])
    return preprocess(client_dataset)
Explanation: We'll use the EMNIST dataset to train a model by loading and preprocessing individual clients (emulating distinct partitions) through an implementation of DataBackend.
Defining a DataBackend
We need an instance of DataBackend to instruct TFF workers (these are the processes that run the client-side of federated computations) how to load and tranform the local data stored in remote locations. A DataBackend is a programmatic construct that resolves symbolic references, represented as application-specific URIs, to materialized payloads that downstream TFF operations can process. Specifically, a DataBackend object is wrapped by a DataExecutor, which queries the object when the TFF runtime encounters an operation that fetches the data.
In this example, an Id to a client is encoded in a URI, which is parsed by our DataBackend definition to retrieve the corresponding client data, convert it to tf.Dataset, and then apply our preprocess function.
End of explanation
def ex_fn(
    device: tf.config.LogicalDevice) -> tff.framework.DataExecutor:
  return tff.framework.DataExecutor(
      tff.framework.EagerTFExecutor(device),
      data_backend=TestDataBackend())
factory = tff.framework.local_executor_factory(leaf_executor_fn=ex_fn)
ctx = tff.framework.ExecutionContext(executor_fn=factory)
tff.framework.set_default_context(ctx)
Explanation: Plugging the DataBackend into the ExecutionContext
TFF computations are invoked by an ExecutionContext and in order for data URIs defined in TFF computations to be understood at runtime, a custom context must be defined that includes a pointer to the DataBackend we just created, so URIs can be properly resolved.
The DataBackend works in tandem with DataExecutor to supply the executor with operable data that the executor can relay to requesting executors in order to complete a TFF computation.
End of explanation
def create_keras_model():
  return tf.keras.models.Sequential([
      tf.keras.layers.InputLayer(input_shape=(784,)),
      tf.keras.layers.Dense(10, kernel_initializer='zeros'),
      tf.keras.layers.Softmax(),
  ])
  
def model_fn():
  keras_model = create_keras_model()
  return tff.learning.from_keras_model(
      keras_model,
      input_spec=preprocessed_example_dataset.element_spec,
      loss=tf.keras.losses.SparseCategoricalCrossentropy(),
      metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])
Explanation: Training the model
Now we are ready to train a model in a federated fashion. Lets define a Keras model along with training hyperparameters:
End of explanation
iterative_process = tff.learning.build_federated_averaging_process(
    model_fn,
    client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.02),
    server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.0))
state = iterative_process.initialize()
Explanation: We can pass this TFF-wrapped definition of our model
to a Federated Averaging algorithm by invoking the helper
function tff.learning.build_federated_averaging_process, as follows:
End of explanation
element_type = tff.types.StructWithPythonType(
    preprocessed_example_dataset.element_spec,
    container_type=collections.OrderedDict)
dataset_type = tff.types.SequenceType(element_type)
data_uris = [f'uri://{i}' for i in range(5)]
data_handle = tff.framework.CreateDataDescriptor(arg_uris=data_uris, arg_type=dataset_type)
Explanation: The initialize computation returns the initial state of the
Federated Averaging process.
To run a round of training, we need to construct a sample of data by organizing a sample of URI references as follows:
End of explanation
state, metrics = iterative_process.next(state, data_handle)
print('round 1, metrics={}'.format(metrics))
Explanation: Now we can round a round of training:
End of explanation
NUM_ROUNDS = 11
for round_num in range(2, NUM_ROUNDS):
  state, metrics = iterative_process.next(state, data_handle)
  print('round {:2d}, metrics={}'.format(round_num, metrics))
Explanation: And we can run a few more rounds:
End of explanation
<END_TASK>