Apprentissage fédéré pour la génération de texte

Voir sur TensorFlow.org Exécuter dans Google Colab Afficher la source sur GitHub Télécharger le cahier

Ce didacticiel s'appuie sur les concepts du didacticiel Federated Learning pour la classification d'images et présente plusieurs autres approches utiles pour l'apprentissage fédéré.

En particulier, nous chargeons un modèle Keras préalablement formé et le raffinons à l'aide d'un entraînement fédéré sur un ensemble de données décentralisé (simulé). Ceci est pratiquement important pour plusieurs raisons. La possibilité d'utiliser des modèles sérialisés facilite le mélange de l'apprentissage fédéré avec d'autres approches de ML. En outre, cela permet d'utiliser une gamme croissante de modèles pré-entraînés - par exemple, la formation de modèles de langage à partir de zéro est rarement nécessaire, car de nombreux modèles pré-entraînés sont maintenant largement disponibles (voir, par exemple, TF Hub ). Au lieu de cela, il est plus judicieux de partir d'un modèle pré-entraîné et de le raffiner à l'aide de Federated Learning, en s'adaptant aux caractéristiques particulières des données décentralisées pour une application particulière.

Pour ce tutoriel, nous commençons avec un RNN qui génère des caractères ASCII, et nous l'affinons via l'apprentissage fédéré. Nous montrons également comment les poids finaux peuvent être renvoyés au modèle Keras d'origine, permettant une évaluation et une génération de texte faciles à l'aide d'outils standard.

!pip install --quiet --upgrade tensorflow-federated-nightly
!pip install --quiet --upgrade nest-asyncio

import nest_asyncio
nest_asyncio.apply()
import collections
import functools
import os
import time

import numpy as np
import tensorflow as tf
import tensorflow_federated as tff

np.random.seed(0)

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

Charger un modèle pré-entraîné

Nous chargeons un modèle qui a été pré-entraîné à la suite du tutoriel TensorFlow Génération de texte à l'aide d'un RNN avec une exécution rapide . Cependant, plutôt que de nous entraîner sur The Complete Works of Shakespeare , nous avons pré-formé le modèle sur le texte de A Tale of Two Cities and A Christmas Carol de Charles Dickens.

À part élargir le vocabulaire, nous n'avons pas modifié le didacticiel original, donc ce modèle initial n'est pas à la pointe de la technologie, mais il produit des prédictions raisonnables et est suffisant pour nos besoins de didacticiel. Le modèle final a été enregistré avec tf.keras.models.save_model(include_optimizer=False) .

Nous utiliserons l'apprentissage fédéré pour affiner ce modèle pour Shakespeare dans ce didacticiel, en utilisant une version fédérée des données fournies par TFF.

Générer les tables de recherche de vocabulaire

# 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)

Chargez le modèle pré-entraîné et générez du texte

def load_model(batch_size):
  urls = {
      1: 'https://storage.googleapis.com/tff-models-public/dickens_rnn.batch1.kerasmodel',
      8: 'https://storage.googleapis.com/tff-models-public/dickens_rnn.batch8.kerasmodel'}
  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)  
  return tf.keras.models.load_model(local_file, compile=False)
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.random.categorical(
        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))
# Text generation requires a batch_size=1 model.
keras_model_batch1 = load_model(batch_size=1)
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
16203776/16193984 [==============================] - 0s 0us/step
What of TensorFlow Federated, you ask? Sall
yesterday. Received the Bailey."

"Mr. Lorry, grimmering himself, or low varked thends the winter, and the eyes of Monsieur
Defarge. "Let his mind, hon in his
life and message; four declare

Charger et prétraiter les données fédérées de Shakespeare

Le package tff.simulation.datasets fournit une variété d'ensembles de données qui sont divisés en «clients», où chaque client correspond à un ensemble de données sur un périphérique particulier qui pourrait participer à l'apprentissage fédéré.

Ces ensembles de données fournissent des distributions de données réalistes non-IID qui reproduisent dans la simulation les défis de la formation sur des données décentralisées réelles. Une partie du prétraitement de ces données a été réalisée à l'aide d'outils du projet Leaf ( github ).

train_data, test_data = tff.simulation.datasets.shakespeare.load_data()

Les jeux de données fournis par shakespeare.load_data() composent d'une séquence de Tensors de chaîne, un pour chaque ligne prononcée par un personnage particulier dans une pièce de Shakespeare. Les clés client sont constituées du nom du jeu joint au nom du personnage, donc par exemple MUCH_ADO_ABOUT_NOTHING_OTHELLO correspond aux lignes du personnage Othello dans la pièce Much Ado About Nothing . Notez que dans un véritable scénario d'apprentissage fédéré, les clients ne sont jamais identifiés ou suivis par des identifiants, mais pour la simulation, il est utile de travailler avec des ensembles de données clés.

Ici, par exemple, nous pouvons regarder quelques données de 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'', shape=(), dtype=string)
tf.Tensor(b'What?', shape=(), dtype=string)

Nous utilisonstf.data.Dataset transformationstf.data.Dataset pour préparer ces données pour l'apprentissage du char RNN chargé ci-dessus.

# Input pre-processing parameters
SEQ_LENGTH = 100
BATCH_SIZE = 8
BUFFER_SIZE = 100  # For dataset shuffling
# Construct a lookup table to map string chars to indexes,
# using the vocab loaded above:
table = tf.lookup.StaticHashTable(
    tf.lookup.KeyValueTensorInitializer(
        keys=vocab, values=tf.constant(list(range(len(vocab))),
                                       dtype=tf.int64)),
    default_value=0)


def to_ids(x):
  s = tf.reshape(x['snippets'], shape=[1])
  chars = tf.strings.bytes_split(s).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 (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
      .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))

Notez que dans la formation des séquences originales et dans la formation des lots ci-dessus, nous utilisons drop_remainder=True pour plus de simplicité. Cela signifie que tous les caractères (clients) qui n'ont pas au moins (SEQ_LENGTH + 1) * BATCH_SIZE chars de texte auront des ensembles de données vides. Une approche typique pour résoudre ce problème serait de remplir les lots avec un jeton spécial, puis de masquer la perte pour ne pas prendre en compte les jetons de remplissage.

Cela compliquerait quelque peu l'exemple, donc pour ce didacticiel, nous n'utilisons que des lots complets, comme dans le didacticiel standard . Cependant, dans le paramètre fédéré, ce problème est plus important, car de nombreux utilisateurs peuvent avoir de petits ensembles de données.

Nous pouvons maintenant prétraiter notre raw_example_dataset et vérifier les types:

example_dataset = preprocess(raw_example_dataset)
print(example_dataset.element_spec)
(TensorSpec(shape=(8, 100), dtype=tf.int64, name=None), TensorSpec(shape=(8, 100), dtype=tf.int64, name=None))

Compilez le modèle et testez sur les données prétraitées

Nous avons chargé un modèle keras non compilé, mais pour exécuter keras_model.evaluate , nous devons le compiler avec une perte et des métriques. Nous compilerons également dans un optimiseur, qui sera utilisé comme optimiseur sur l'appareil dans Federated Learning.

Le didacticiel d'origine n'avait pas de précision au niveau des caractères (la fraction des prédictions où la probabilité la plus élevée était placée sur le caractère suivant correct). C'est une métrique utile, nous l'ajoutons donc. Cependant, nous devons définir une nouvelle classe de métrique pour cela car nos prédictions ont le rang 3 (un vecteur de logits pour chacune des prédictions BATCH_SIZE * SEQ_LENGTH ), et SparseCategoricalAccuracy n'attend que des prédictions de rang 2.

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

  def __init__(self, name='accuracy', dtype=tf.float32):
    super().__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().update_state(y_true, y_pred, sample_weight)

Nous pouvons maintenant compiler un modèle et l'évaluer sur notre example_dataset .

BATCH_SIZE = 8  # The training and eval batch size for the rest of this tutorial.
keras_model = load_model(batch_size=BATCH_SIZE)
keras_model.compile(
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=[FlattenedCategoricalAccuracy()])

# Confirm that loss is much lower on Shakespeare than on random data
loss, accuracy = keras_model.evaluate(example_dataset.take(5), verbose=0)
print(
    'Evaluating on an example Shakespeare character: {a:3f}'.format(a=accuracy))

# As a sanity check, we can construct some completely random data, where we expect
# the accuracy to be essentially random:
random_guessed_accuracy = 1.0 / len(vocab)
print('Expected accuracy for random guessing: {a:.3f}'.format(
    a=random_guessed_accuracy))
random_indexes = np.random.randint(
    low=0, high=len(vocab), size=1 * BATCH_SIZE * (SEQ_LENGTH + 1))
data = collections.OrderedDict(
    snippets=tf.constant(
        ''.join(np.array(vocab)[random_indexes]), shape=[1, 1]))
random_dataset = preprocess(tf.data.Dataset.from_tensor_slices(data))
loss, accuracy = keras_model.evaluate(random_dataset, steps=10, verbose=0)
print('Evaluating on completely random data: {a:.3f}'.format(a=accuracy))
Downloading data from https://storage.googleapis.com/tff-models-public/dickens_rnn.batch8.kerasmodel
16195584/16193984 [==============================] - 0s 0us/step
16203776/16193984 [==============================] - 0s 0us/step
Evaluating on an example Shakespeare character: 0.402000
Expected accuracy for random guessing: 0.012
Evaluating on completely random data: 0.011

Affinez le modèle avec Federated Learning

TFF sérialise tous les calculs TensorFlow afin qu'ils puissent potentiellement être exécutés dans un environnement non Python (même si pour le moment, seul un runtime de simulation implémenté en Python est disponible). Même si nous fonctionnons en mode impatient (TF 2.0), actuellement TFF sérialise les calculs TensorFlow en construisant les opérations nécessaires dans le contexte d'une with tf.Graph.as_default() " with tf.Graph.as_default() ". Ainsi, nous devons fournir une fonction que TFF peut utiliser pour introduire notre modèle dans un graphe qu'il contrôle. Nous procédons comme suit:

# 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. Note: we want to construct all the necessary objects we'll need 
# _inside_ this method.
def create_tff_model():
  # TFF uses an `input_spec` so it knows the types and shapes
  # that your model expects.
  input_spec = example_dataset.element_spec
  keras_model_clone = tf.keras.models.clone_model(keras_model)
  return tff.learning.from_keras_model(
      keras_model_clone,
      input_spec=input_spec,
      loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
      metrics=[FlattenedCategoricalAccuracy()])

Nous sommes maintenant prêts à construire un processus itératif de moyenne fédérée, que nous utiliserons pour améliorer le modèle (pour plus de détails sur l'algorithme de moyenne fédérée, voir l'étude Communication-Efficient Learning of Deep Networks from Decentralized Data ).

Nous utilisons un modèle Keras compilé pour effectuer une évaluation standard (non fédérée) après chaque cycle de formation fédérée. Ceci est utile à des fins de recherche lors de l'apprentissage fédéré simulé et il existe un jeu de données de test standard.

Dans un contexte de production réaliste, cette même technique peut être utilisée pour prendre des modèles formés avec un apprentissage fédéré et les évaluer sur un ensemble de données de référence centralisé à des fins de test ou d'assurance qualité.

# This command builds all the TensorFlow graphs and serializes them: 
fed_avg = tff.learning.build_federated_averaging_process(
    model_fn=create_tff_model,
    client_optimizer_fn=lambda: tf.keras.optimizers.SGD(lr=0.5))

Voici la boucle la plus simple possible, dans laquelle nous exécutons une moyenne fédérée pour un tour sur un seul client sur un seul lot:

state = fed_avg.initialize()
state, metrics = fed_avg.next(state, [example_dataset.take(5)])
train_metrics = metrics['train']
print('loss={l:.3f}, accuracy={a:.3f}'.format(
    l=train_metrics['loss'], a=train_metrics['accuracy']))
loss=4.403, accuracy=0.132

Écrivons maintenant une boucle d'entraînement et d'évaluation un peu plus intéressante.

Pour que cette simulation fonctionne encore relativement rapidement, nous nous entraînons sur les trois mêmes clients à chaque tour, en ne considérant que deux minibatchs pour chacun.

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


clients = [
    'ALL_S_WELL_THAT_ENDS_WELL_CELIA', 'MUCH_ADO_ABOUT_NOTHING_OTHELLO',
]

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

# We concatenate the test datasets for evaluation with Keras by creating a 
# Dataset of Datasets, and then identity flat mapping across all the examples.
test_dataset = tf.data.Dataset.from_tensor_slices(
    [data(client, test_data) for client in clients]).flat_map(lambda x: x)

L'état initial du modèle produit par fed_avg.initialize() est basé sur les initialiseurs aléatoires du modèle Keras, pas sur les poids chargés, puisque clone_model() ne clone pas les poids. Pour démarrer l'entraînement à partir d'un modèle pré-entraîné, nous définissons les pondérations du modèle dans l'état du serveur directement à partir du modèle chargé.

NUM_ROUNDS = 5

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

# Load our pre-trained Keras model weights into the global model state.
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):
  # Take our global model weights and push them back into a Keras model to
  # use its standard `.evaluate()` method.
  keras_model = load_model(batch_size=BATCH_SIZE)
  keras_model.compile(
      loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
      metrics=[FlattenedCategoricalAccuracy()])
  state.model.assign_weights_to(keras_model)
  loss, accuracy = keras_model.evaluate(example_dataset, steps=2, verbose=0)
  print('\tEval: loss={l:.3f}, accuracy={a:.3f}'.format(l=loss, a=accuracy))


for round_num in range(NUM_ROUNDS):
  print('Round {r}'.format(r=round_num))
  keras_evaluate(state, round_num)
  state, metrics = fed_avg.next(state, train_datasets)
  train_metrics = metrics['train']
  print('\tTrain: loss={l:.3f}, accuracy={a:.3f}'.format(
      l=train_metrics['loss'], a=train_metrics['accuracy']))

print('Final evaluation')
keras_evaluate(state, NUM_ROUNDS + 1)
Round 0
    Eval: loss=3.324, accuracy=0.401
    Train: loss=4.360, accuracy=0.155
Round 1
    Eval: loss=4.361, accuracy=0.049
    Train: loss=4.235, accuracy=0.164
Round 2
    Eval: loss=4.219, accuracy=0.177
    Train: loss=4.081, accuracy=0.221
Round 3
    Eval: loss=4.080, accuracy=0.174
    Train: loss=3.940, accuracy=0.226
Round 4
    Eval: loss=3.991, accuracy=0.176
    Train: loss=3.840, accuracy=0.226
Final evaluation
    Eval: loss=3.909, accuracy=0.171

Avec les modifications par défaut, nous n'avons pas fait suffisamment de formation pour faire une grande différence, mais si vous vous entraînez plus longtemps sur plus de données Shakespeare, vous devriez voir une différence dans le style du texte généré avec le modèle mis à jour:

# Set our newly trained weights back in the originally created model.
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? Shalways, I will call your
compet with any city brought their faces uncompany," besumed him. "When he
sticked Madame Defarge pushed the lamps.

"Have I often but no unison. She had probably come,

Extensions suggérées

Ce tutoriel n'est que la première étape! Voici quelques idées sur la façon dont vous pourriez essayer d'étendre ce bloc-notes:

  • Rédigez une boucle de formation plus réaliste dans laquelle vous échantillonnez les clients sur lesquels vous vous entraînez de manière aléatoire.
  • Utilisez " .repeat(NUM_EPOCHS) " sur les ensembles de données client pour essayer plusieurs époques de formation locale (par exemple, comme dans McMahan et. Al. ). Voir également Apprentissage fédéré pour la classification d'images qui effectue cette opération.
  • Modifiez la commande compile() pour expérimenter l'utilisation de différents algorithmes d'optimisation sur le client.
  • Essayez l'argument server_optimizer à build_federated_averaging_process pour essayer différents algorithmes pour appliquer les mises à jour de modèle sur le serveur.
  • Essayez l'argument client_weight_fn à pour build_federated_averaging_process pour essayer différentes pondérations des clients. La valeur par défaut pondère les mises à jour du client en fonction du nombre d'exemples sur le client, mais vous pouvez faire par exemple client_weight_fn=lambda _: tf.constant(1.0) .