TensorFlow 2.0 Beta is available

# Federated Learning for Text Generation

NOTE: This colab has been verified to work with the 0.6.0 version of the tensorflow_federated pip package, but the Tensorflow Federated project is still in pre-release development and may not work on master.

This tutorial builds on the concepts in the Federated Learning for Image Classification tutorial, and demonstrates several other useful approaches for federated learning.

In particular, we load a previously trained Keras model, and refine it using federated training on a (simulated) decentralized dataset. This is practically important for several reasons . The ability to use serialized models makes it easy to mix federated learning with other ML approaches. Further, this allows use of an increasing range of pre-trained models --- for example, training language models from scratch is rarely necessary, as numerous pre-trained models are now widely available (see, e.g., TF Hub). Instead, it makes more sense to start from a pre-trained model, and refine it using Federated Learning, adapting to the particular characteristics of the decentralized data for a particular application.

For this tutorial, we start with a RNN that generates ASCII characters, and refine it via federated learning. We also show how the final weights can be fed back to the original Keras model, allowing easy evaluation and text generation using standard tools.

#@test {"skip": true}

# NOTE: If you are running a Jupyter notebook, and installing a locally built
# pip package, you may need to edit the following to point to the '.whl' file

!pip install --quiet tensorflow_federated
!pip install --quiet  tf-nightly

DEPRECATION: Python 3.4 support has been deprecated. pip 19.1 will be the last one supporting it. Please upgrade your Python as Python 3.4 won't be maintained after March 2019 (cf PEP 429).
DEPRECATION: Python 3.4 support has been deprecated. pip 19.1 will be the last one supporting it. Please upgrade your Python as Python 3.4 won't be maintained after March 2019 (cf PEP 429).

from __future__ import absolute_import, division, print_function

import collections
import functools
import os
import time

import numpy as np
import tensorflow as tf

import tensorflow_federated as tff

tf.compat.v1.enable_v2_behavior()

np.random.seed(0)

# Test the TFF is working:
tff.federated_computation(lambda: 'Hello, World!')()

WARNING: Logging before flag parsing goes to stderr.
The TensorFlow contrib module will not be included in TensorFlow 2.0.
* https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md
* https://github.com/tensorflow/io (for I/O related ops)
If you depend on functionality not listed there, please file an issue.

'Hello, World!'

We load a model that was pre-trained following the TensorFlow tutorial Text generation using a RNN with eager execution. However, rather than training on The Complete Works of Shakespeare, we pre-trained the model on the text from the Charles Dickens' A Tale of Two Cities and A Christmas Carol.

Other than expanding the vocabularly, we didn't modify the original tutorial, so this initial model isn't state-of-the-art, but it produces reasonable predictions and is sufficient for our tutorial purposes. The final model was saved with tf.keras.models.save_model(include_optimizer=False).

We will use federated learning to fine-tune this model for Shakespeare in this tutorial, using a federated version of the data provided by TFF.

## Generate the vocab lookup tables

# A fixed vocabularly of ASCII chars that occur in the works of Shakespeare and Dickens:
vocab = list('dhlptx@DHLPTX \$(,048cgkoswCGKOSW[_#\'/37;?bfjnrvzBFJNRVZ"&*.26:\naeimquyAEIMQUY]!%)-159\r')

# Creating a mapping from unique characters to indices
char2idx = {u:i for i, u in enumerate(vocab)}
idx2char = np.array(vocab)


## Load the pre-trained model and generate some text

def load_model(batch_size):
urls = {
assert batch_size in urls, 'batch_size must be in ' + str(urls.keys())
url = urls[batch_size]
local_file = tf.keras.utils.get_file(os.path.basename(url), origin=url)

def generate_text(model, start_string):
# From https://www.tensorflow.org/tutorials/sequences/text_generation
num_generate = 200
input_eval = [char2idx[s] for s in start_string]
input_eval = tf.expand_dims(input_eval, 0)
text_generated = []
temperature = 1.0

model.reset_states()
for i in range(num_generate):
predictions = model(input_eval)
predictions = tf.squeeze(predictions, 0)
predictions = predictions / temperature
predicted_id = tf.multinomial(predictions, num_samples=1)[-1,0].numpy()
input_eval = tf.expand_dims([predicted_id], 0)
text_generated.append(idx2char[predicted_id])

return (start_string + ''.join(text_generated))

#@test {"output": "ignore", "timeout": 120}
# Text generation requires a batch_size=1 model.
print(generate_text(keras_model_batch1, 'What of TensorFlow Federated, you ask? '))

Downloading data from https://storage.googleapis.com/tff-models-public/dickens_rnn.batch1.kerasmodel
16195584/16193984 [==============================] - 0s 0us/step

W0626 14:54:38.861224 140226356918080 deprecation.py:323] From <ipython-input-6-5df7f7779638>:14: multinomial (from tensorflow.python.ops.random_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use tf.random.categorical instead.

What of TensorFlow Federated, you ask? Shrewn upon the prisoner, "though it is not so
hot," the great figure been found in a pig, and the pulle
of wine for Liberty, Equality, Fraternity, or Death,
Monseigneur was when they rested on the


# Load and Preprocess the Federated Shakespeare Data

The tff.simulation.datasets package provides a variety of datasets that are split into "clients", where each client corresponds to a dataset on a particular device that might participate in federated learning.

These datasets provide realistic non-IID data distributions that replicate in simulation the challenges of training on real decentralized data. Some of the pre-processing of this data was done using tools from the Leaf project (github).

#@test {"output": "ignore", "timeout": 120}

Downloading data from https://storage.googleapis.com/tff-datasets-public/shakespeare.tar.bz2
1851392/1848122 [==============================] - 0s 0us/step

W0626 14:54:42.690933 140226356918080 deprecation.py:323] From /tmpfs/src/tf_docs_env/lib/python3.4/site-packages/tensorflow_federated/python/simulation/hdf5_client_data.py:64: DatasetV1.output_types (from tensorflow.python.data.ops.dataset_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use tf.compat.v1.data.get_output_types(dataset).
W0626 14:54:42.692488 140226356918080 deprecation.py:323] From /tmpfs/src/tf_docs_env/lib/python3.4/site-packages/tensorflow_federated/python/simulation/hdf5_client_data.py:65: DatasetV1.output_shapes (from tensorflow.python.data.ops.dataset_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use tf.compat.v1.data.get_output_shapes(dataset).


The datasets provided by shakespeare.load_data() consist of a sequence of string Tensors, one for each line spoken by a particular character in a Shakespeare play. The client keys consist of the name of the play joined with the name of the character, so for example MUCH_ADO_ABOUT_NOTHING_OTHELLO corresponds to the lines for the character Othello in the play Much Ado About Nothing. Note that in a real federated learning scenario clients are never identified or tracked by ids, but for simulation it is useful to work with keyed datasets.

Here, for example, we can look at some data from King Lear:

# Here the play is "The Tragedy of King Lear" and the character is "King".
raw_example_dataset = train_data.create_tf_dataset_for_client(
'THE_TRAGEDY_OF_KING_LEAR_KING')
# To allow for future extensions, each entry x
# is an OrderedDict with a single key 'snippets' which contains the text.
for x in raw_example_dataset.take(2):
print(x['snippets'])

tf.Tensor(b"Live regist'red upon our brazen tombs,\nAnd then grace us in the disgrace of death;\nWhen, spite of cormorant devouring Time,\nTh' endeavour of this present breath may buy\nThat honour which shall bate his scythe's keen edge,\nAnd make us heirs of all eternity.\nTherefore, brave conquerors- for so you are\nThat war against your own affections\nAnd the huge army of the world's desires-\nOur late edict shall strongly stand in force:\nNavarre shall be the wonder of the world;\nOur court shall be a little Academe,\nStill and contemplative in living art.\nYou three, Berowne, Dumain, and Longaville,\nHave sworn for three years' term to live with me\nMy fellow-scholars, and to keep those statutes\nThat are recorded in this schedule here.\nYour oaths are pass'd; and now subscribe your names,\nThat his own hand may strike his honour down\nThat violates the smallest branch herein.\nIf you are arm'd to do as sworn to do,\nSubscribe to your deep oaths, and keep it too.\nYour oath is pass'd to pass away from these.", shape=(), dtype=string)
tf.Tensor(b'this?\nDid you hear the proclamation?', shape=(), dtype=string)


We now use tf.data.Dataset transformations to prepare this data for training the char RNN loaded above.

# Input pre-processing parameters
SEQ_LENGTH = 100
BATCH_SIZE = 8
BUFFER_SIZE = 10000  # For dataset shuffling

# Using a namedtuple with keys x and y as the output type of the
# dataset keeps both TFF and Keras happy:
BatchType = collections.namedtuple('BatchType', ['x', 'y'])

# Construct a lookup table to map string chars to indexes,
# using the vocab loaded above:
table = tf.contrib.lookup.index_table_from_tensor(
mapping=vocab,
num_oov_buckets=0,
default_value=0)

def to_ids(x):
s = tf.reshape(x['snippets'], shape=[1])
chars = tf.string_split(s, delimiter='').values
ids = table.lookup(chars)
return ids

def split_input_target(chunk):
input_text = tf.map_fn(lambda x: x[:-1], chunk)
target_text = tf.map_fn(lambda x: x[1:], chunk)
return BatchType(input_text, target_text)

def preprocess(dataset):
return (
# Map ASCII chars to int64 indexes using the vocab
dataset.map(to_ids)
# Split into individual chars
.apply(tf.data.experimental.unbatch())
# Form example sequences of SEQ_LENGTH +1
.batch(SEQ_LENGTH + 1,  drop_remainder=True)
# Shuffle and form minibatches
.shuffle(BUFFER_SIZE).batch(BATCH_SIZE, drop_remainder=True)
# And finally split into (input, target) tuples,
# each of length SEQ_LENGTH.
.map(split_input_target))


Note that in the formation of the original sequences and in the formation of batches above, we use drop_remainder=True for simplicity. This means that any characters (clients) that don't have at least (SEQ_LENGTH + 1) * BATCH_SIZE chars of text will have empty datasets. A typical approach to address this would be to pad the batches with a special token, and then mask the loss to not take the padding tokens into account.

This would complicate the example somewhat, so for this tutorial we only use full batches, as in the standard tutorial. However, in the federated setting this issue is more significant, because many users might have small datasets.

Now we can preprocess our raw_example_dataset, and check the types:

#@test {"output": "ignore"}
example_dataset = preprocess(raw_example_dataset)
print(example_dataset.output_types, example_dataset.output_shapes)

W0626 14:54:42.794631 140226356918080 deprecation_wrapper.py:118] From /home/kbuilder/.local/lib/python3.4/site-packages/tensorflow_estimator/python/estimator/api/_v1/estimator/__init__.py:10: The name tf.estimator.inputs is deprecated. Please use tf.compat.v1.estimator.inputs instead.

W0626 14:54:42.899141 140226356918080 deprecation.py:506] From /tmpfs/src/tf_docs_env/lib/python3.4/site-packages/tensorflow_core/python/autograph/impl/api.py:350: calling string_split (from tensorflow.python.ops.ragged.ragged_string_ops) with delimiter is deprecated and will be removed in a future version.
Instructions for updating:
W0626 14:54:42.907609 140226356918080 deprecation.py:323] From <ipython-input-11-ac7fca9ee277>:29: unbatch (from tensorflow.python.data.experimental.ops.batching) is deprecated and will be removed in a future version.
Instructions for updating:
Use tf.data.Dataset.unbatch().

BatchType(x=tf.int64, y=tf.int64) BatchType(x=TensorShape([8, 100]), y=TensorShape([8, 100]))


# Compile the model and test on the preprocessed data

We loaded an uncompiled keras model, but in order to run keras_model.evaluate, we need to compile it with a loss and metrics. We will also compile in an optimizer, which will be used as the on-device optimizer in Federated Learning.

The original tutorial didn't have char-level accuracy (the fraction of predictions where the highest probability was put on the correct next char). This is a useful metric, so we add it. However, we need to define a new metric class for this because our predictions have rank 3 (a vector of logits for each of the BATCH_SIZE * SEQ_LENGTH predictions), and SparseCategoricalAccuracy expects only rank 2 predictions.

class FlattenedCategoricalAccuracy(tf.keras.metrics.SparseCategoricalAccuracy):

def __init__(self, name='accuracy', dtype=None):
super(FlattenedCategoricalAccuracy, self).__init__(name, dtype=dtype)

def update_state(self, y_true, y_pred, sample_weight=None):
y_true = tf.reshape(y_true, [-1, 1])
y_pred = tf.reshape(y_pred, [-1, len(vocab), 1])
return super(FlattenedCategoricalAccuracy, self).update_state(
y_true, y_pred, sample_weight)

def compile(keras_model):
keras_model.compile(
optimizer=tf.keras.optimizers.SGD(lr=0.5),
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=[FlattenedCategoricalAccuracy()])
return keras_model


Now we can compile a model, and evaluate it on our example_dataset.

#@test {"output": "ignore", "timeout": 120}
BATCH_SIZE = 8  # The training and eval batch size for the rest of this tutorial.

compile(keras_model)

# Confirm that loss is much lower on Shakespeare than on random data
print('Evaluating on an example Shakespeare character:')
keras_model.evaluate(example_dataset.take(1))

# As a sanity check, we can construct some completely random data, where we expect
# the accuracy to be essentially random:
random_indexes = np.random.randint(
low=0, high=len(vocab), size=1 * BATCH_SIZE * (SEQ_LENGTH + 1))
data = {
'snippets':
tf.constant(''.join(np.array(vocab)[random_indexes]), shape=[1, 1])
}
random_dataset = preprocess(tf.data.Dataset.from_tensor_slices(data))
print('Expected accuracy for random guessing: {:.3f}'.format(1.0 / len(vocab)))
print('Evaluating on completely random data:')
keras_model.evaluate(random_dataset, steps=1)

Downloading data from https://storage.googleapis.com/tff-models-public/dickens_rnn.batch8.kerasmodel
16195584/16193984 [==============================] - 0s 0us/step
Evaluating on an example Shakespeare character:
1/Unknown - 1s 822ms/step - loss: 3.2307 - accuracy: 0.3963Expected accuracy for random guessing: 0.012
Evaluating on completely random data:
1/1 [==============================] - 1s 568ms/step - loss: 11.4394 - accuracy: 0.0125

[11.439443588256836, 0.0125]

# Fine-tune the model with Federated Learning

TFF serializes all TensorFlow computations so they can potentially be run in a non-Python environment (even though at the moment, only a simulation runtime implemented in Python is available). Even though we are running in eager mode, (TF 2.0), currently TFF serializes TensorFlow computations by constructing the necessary ops inside the context of a "with tf.Graph.as_default()" statement. Thus, we need to provide a function that TFF can use to introduce our model into a graph it controls. We do this as follows:

# Clone the keras_model inside create_tff_model(), which TFF will
# call to produce a new copy of the model inside the graph that it will serialize.
def create_tff_model():
# TFF uses a dummy_batch so it knows the types and shapes
x = tf.constant(np.random.randint(1, len(vocab), size=[BATCH_SIZE, SEQ_LENGTH]))
dummy_batch = collections.OrderedDict([('x', x), ('y', x)])
keras_model_clone = compile(tf.keras.models.clone_model(keras_model))
return tff.learning.from_compiled_keras_model(
keras_model_clone, dummy_batch=dummy_batch)


Now we are ready to construct a Federated Averaging iterative process, which we will use to improve the model (for details on the Federated Averaging algorithm, see the paper Communication-Efficient Learning of Deep Networks from Decentralized Data).

We use a compiled Keras model to perform standard (non-federated) evaluation after each round of federated training. This is useful for research purposes when doing simulated federated learning and there is a standard test dataset.

In a realistic production setting this same technique might be used to take models trained with federated learning and evaluate them on a centralized benchmark dataset for testing or quality assurance purposes.

#@test {"output": "ignore", "timeout": 300}
# This command builds all the TensorFlow graphs and serializes them:
fed_avg = tff.learning.build_federated_averaging_process(model_fn=create_tff_model)

W0626 14:54:45.495552 140226356918080 deprecation.py:506] From /tmpfs/src/tf_docs_env/lib/python3.4/site-packages/tensorflow_core/python/ops/resource_variable_ops.py:1624: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.
Instructions for updating:
If using Keras pass *_constraint arguments to layers.

W0626 14:54:49.178100 140226356918080 deprecation.py:323] From /tmpfs/src/tf_docs_env/lib/python3.4/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py:460: BaseResourceVariable.constraint (from tensorflow.python.ops.resource_variable_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Apply a constraint manually following the optimizer update step.


Here is the simplest possible loop, where we run federated averaging for one round on a single client on a single batch:

#@test {"output": "ignore", "timeout": 600}
state = fed_avg.initialize()
state, metrics = fed_avg.next(state, [example_dataset.take(1)])
print(metrics)

<accuracy=0.00875,loss=4.454028>


Now let's write a slightly more interesting training and evaluation loop.

So that this simulation still runs relatively quickly, we train on the same three clients each round, only considering two minibatches for each.

def data(client, source=train_data):
return preprocess(
source.create_tf_dataset_for_client(client)).take(2)

clients = ['ALL_S_WELL_THAT_ENDS_WELL_CELIA',
'THE_TRAGEDY_OF_KING_LEAR_KING']

train_datasets = [data(client) for client in clients]

# We concatenate the test datasets for evaluation with Keras.
test_dataset = functools.reduce(
lambda d1, d2: d1.concatenate(d2),
[data(client, test_data) for client in clients])


The initial state of the model produced by fed_avg.initialize() is based on the random initializers for the Keras model, not the weights that were loaded, since clone_model() does not clone the weights. To start training from a pre-trained model, we set the model weights in the server state directly from the loaded model.

#@test {"skip": true}
NUM_ROUNDS = 3

# The state of the FL server, containing the model and optimization state.
state = fed_avg.initialize()

state = tff.learning.state_with_new_model_weights(
state,
trainable_weights=[v.numpy() for v in keras_model.trainable_weights],
non_trainable_weights=[
v.numpy() for v in keras_model.non_trainable_weights
])

def keras_evaluate(state, round_num):
tff.learning.assign_weights_to_keras_model(keras_model, state.model)
print('Evaluating before training round', round_num)
keras_model.evaluate(example_dataset, steps=2)

for round_num in range(NUM_ROUNDS):
keras_evaluate(state, round_num)
# N.B. The TFF runtime is currently fairly slow,
# expect this to get significantly faster in future releases.
state, metrics = fed_avg.next(state, train_datasets)
print('Training metrics: ', metrics)

keras_evaluate(state, NUM_ROUNDS + 1)

Evaluating before training round 0
2/2 [==============================] - 1s 580ms/step - loss: 3.3203 - accuracy: 0.3938
Training metrics:  <accuracy=0.41875,loss=3.1005116>
Evaluating before training round 1
2/2 [==============================] - 1s 481ms/step - loss: 2.9440 - accuracy: 0.4269
Training metrics:  <accuracy=0.53125,loss=2.2019594>
Evaluating before training round 2
2/2 [==============================] - 1s 589ms/step - loss: 2.7216 - accuracy: 0.4512
Training metrics:  <accuracy=0.6325,loss=1.5981749>
Evaluating before training round 4
2/2 [==============================] - 1s 525ms/step - loss: 2.6187 - accuracy: 0.4669


With the default changes, we haven't done enough training to make a big difference, but if you train longer on more Shakespeare data, you should see a difference in the style of the text generated with the updated model:

#@test {"output": "ignore", "timeout": 300}
keras_model_batch1.set_weights([v.numpy() for v in keras_model.weights])
# Text generation requires batch_size=1
print(generate_text(keras_model_batch1, 'What of TensorFlow Federated, you ask? '))

What of TensorFlow Federated, you ask? Why soothing patriot says, all. Since then, brother?'

"There I say y to Mr. Lorry't change. It was haught.
Alourged."

"Found."

"None, sir, you a gentlemen,'d natural ways; of those thousands


# Suggested extensions

This tutorial is just the first step! Here are some ideas for how you might try extending this notebook: * Write a more realistic training loop where you sample clients to train on randomly. * Use ".repeat(NUM_EPOCHS)" on the client datasets to try multiple epochs of local training (e.g., as in McMahan et. al.). See also Federated Learning for Image Classification which does this. * Change the compile() command to experiment with using different optimization algorithms on the client. * Try the server_optimizer argument to build_federated_averaging_process to try different algorithms for applying the model updates on the server. * Try the client_weight_fn argument to to build_federated_averaging_process to try different weightings of the clients. The default weights client updates by the number of examples on the client, but you can do e.g. client_weight_fn=lambda _: tf.constant(1.0).