Remove the dependency from sample code to examples.
@@ -52,31 +52,34 @@ landing_page:
and try it out yourself!
code_block: |
<pre class = "prettyprint">
- from six.moves import range
import tensorflow as tf
- import tensorflow_federated as tff
- from tensorflow_federated.python.examples import mnist
tf.compat.v1.enable_v2_behavior()
+ import tensorflow_federated as tff
- # Load simulation data.
source, _ = tff.simulation.datasets.emnist.load_data()
def client_data(n):
- dataset = source.create_tf_dataset_for_client(source.client_ids[n])
- return mnist.keras_dataset_from_emnist(dataset).repeat(10).batch(20)
+ return source.create_tf_dataset_for_client(source.client_ids[n]).map(
+ lambda e: {
+ 'x': tf.reshape(e['pixels'], [-1]),
+ 'y': e['label'],
+ }).repeat(10).batch(20)
- # Pick a subset of client devices to participate in training.
train_data = [client_data(n) for n in range(3)]
- # Grab a single batch of data so that TFF knows what data looks like.
sample_batch = tf.nest.map_structure(
lambda x: x.numpy(), iter(train_data[0]).next())
- # Wrap a Keras model for use with TFF.
def model_fn():
- return tff.learning.from_compiled_keras_model(
- mnist.create_simple_keras_model(), sample_batch)
+ model = tf.keras.models.Sequential([
+ tf.keras.layers.Dense(10, tf.nn.softmax, input_shape=(784,),
+ kernel_initializer='zeros')
+ ])
+ model.compile(
+ loss=tf.keras.losses.SparseCategoricalCrossentropy(),
+ optimizer=tf.keras.optimizers.SGD(0.1),
+ metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])
+ return tff.learning.from_compiled_keras_model(model, sample_batch)
- # Simulate a few rounds of training with the selected client devices.
trainer = tff.learning.build_federated_averaging_process(model_fn)
state = trainer.initialize()
for _ in range(5):
