Simulation haute performance avec Kubernetes

Ce didacticiel décrira comment configurer une simulation haute performance à l'aide d'un environnement d'exécution TFF exécuté sur Kubernetes. Le modèle est le même que dans le précédent tutoriel, simulations haute performance avec TFF. La seule différence est qu'ici nous utilisons un pool de travailleurs au lieu d'un exécuteur local.

Ce tutoriel fait référence à de Google Cloud GKE pour créer le cluster Kubernetes, mais toutes les étapes après la création du cluster peut être utilisé avec toute installation Kubernetes.

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

Lancer les nœuds de calcul TFF sur GKE

Créer un cluster Kubernetes

L'étape suivante n'est à faire qu'une seule fois. Le cluster peut être réutilisé pour de futures charges de travail.

Suivez les instructions de GKE pour créer un cluster de conteneurs . Le reste de ce tutoriel suppose que le cluster est nommé tff-cluster , mais le nom réel est pas important. Ne plus suivre les instructions quand vous arrivez à « Étape 5: Déployer votre application ».

Déployer l'application de travail TFF

Les commandes pour interagir avec les GCP peuvent être exécutées localement ou dans le Google Cloud Shell . Nous recommandons Google Cloud Shell car il ne nécessite aucune configuration supplémentaire.

  1. Exécutez la commande suivante pour lancer l'application Kubernetes.
$ kubectl create deployment tff-workers --image=gcr.io/tensorflow-federated/remote-executor-service:latest
  1. Ajoutez un équilibreur de charge pour l'application.
$ kubectl expose deployment tff-workers --type=LoadBalancer --port 80 --target-port 8000

Recherchez l'adresse IP de l'équilibreur de charge sur Google Cloud Console. Vous en aurez besoin plus tard pour connecter la boucle d'entraînement à l'application Worker.

(Alternativement) Lancer le conteneur Docker localement

$ docker run --rm -p 8000:8000 gcr.io/tensorflow-federated/remote-executor-service:latest

Configurer l'environnement TFF

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

import nest_asyncio
nest_asyncio.apply()

Définir le modèle à former

import collections
import time

import tensorflow as tf
import tensorflow_federated as tff

source, _ = tff.simulation.datasets.emnist.load_data()


def map_fn(example):
  return collections.OrderedDict(
      x=tf.reshape(example['pixels'], [-1, 784]), y=example['label'])


def client_data(n):
  ds = source.create_tf_dataset_for_client(source.client_ids[n])
  return ds.repeat(10).batch(20).map(map_fn)


train_data = [client_data(n) for n in range(10)]
input_spec = train_data[0].element_spec


def model_fn():
  model = tf.keras.models.Sequential([
      tf.keras.layers.InputLayer(input_shape=(784,)),
      tf.keras.layers.Dense(units=10, kernel_initializer='zeros'),
      tf.keras.layers.Softmax(),
  ])
  return tff.learning.from_keras_model(
      model,
      input_spec=input_spec,
      loss=tf.keras.losses.SparseCategoricalCrossentropy(),
      metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])


trainer = tff.learning.build_federated_averaging_process(
    model_fn, client_optimizer_fn=lambda: tf.keras.optimizers.SGD(0.02))


def evaluate(num_rounds=10):
  state = trainer.initialize()
  for round in range(num_rounds):
    t1 = time.time()
    state, metrics = trainer.next(state, train_data)
    t2 = time.time()
    print('Round {}: loss {}, round time {}'.format(round, metrics.loss, t2 - t1))

Configurer les exécuteurs distants

Par défaut, TFF exécute tous les calculs localement. Dans cette étape, nous demandons à TFF de se connecter aux services Kubernetes que nous avons configurés ci-dessus. Assurez-vous de copier l'adresse IP de votre service ici.

import grpc

ip_address = '0.0.0.0' 
port = 80 

channels = [grpc.insecure_channel(f'{ip_address}:{port}') for _ in range(10)]

tff.backends.native.set_remote_execution_context(channels)

Exécuter la formation

evaluate()

Round 0: loss 4.370407581329346, round time 4.201097726821899
Round 1: loss 4.1407670974731445, round time 3.3283166885375977
Round 2: loss 3.865147590637207, round time 3.098310947418213
Round 3: loss 3.534019708633423, round time 3.1565616130828857
Round 4: loss 3.272688388824463, round time 3.175067663192749
Round 5: loss 2.935391664505005, round time 3.008434534072876
Round 6: loss 2.7399251461029053, round time 3.31435227394104
Round 7: loss 2.5054931640625, round time 3.4411356449127197
Round 8: loss 2.290508985519409, round time 3.158798933029175
Round 9: loss 2.1194536685943604, round time 3.1348156929016113