Simulazione ad alte prestazioni con Kubernetes

Questo tutorial descriverà come configurare una simulazione ad alte prestazioni usando un runtime TFF in esecuzione su Kubernetes. Il modello è lo stesso come nel precedente tutorial, simulazioni ad alte prestazioni con TFF. L'unica differenza è che qui utilizziamo un pool di lavoro anziché un esecutore locale.

Questo tutorial si riferisce a di Google Cloud GKE per creare il cluster kubernetes, ma tutti i passi dopo il cluster viene creato può essere utilizzato con qualsiasi installazione kubernetes.

Visualizza su TensorFlow.org Esegui in Google Colab Visualizza la fonte su GitHub Scarica taccuino

Avvia i lavoratori TFF su GKE

Crea un cluster Kubernetes

Il passaggio successivo deve essere eseguito solo una volta. Il cluster può essere riutilizzato per futuri carichi di lavoro.

Seguire le istruzioni GKE per creare un cluster contenitore . Il resto di questo tutorial presuppone che il cluster è denominato tff-cluster , ma il vero nome non è importante. Smettere di seguire le istruzioni quando si arriva a "Punto 5: distribuire l'applicazione".

Distribuire l'applicazione di lavoro TFF

I comandi per interagire con i GCP può essere eseguito localmente o nella cloud di Google Shell . Consigliamo Google Cloud Shell poiché non richiede una configurazione aggiuntiva.

  1. Esegui il comando seguente per avviare l'applicazione Kubernetes.
$ kubectl create deployment tff-workers --image=gcr.io/tensorflow-federated/remote-executor-service:latest
  1. Aggiungi un servizio di bilanciamento del carico per l'applicazione.
$ kubectl expose deployment tff-workers --type=LoadBalancer --port 80 --target-port 8000

Cerca l'indirizzo IP del loadbalancer su Google Cloud Console. Ti servirà in seguito per connettere il ciclo di formazione all'app lavoratore.

(In alternativa) Avvia il Docker Container in locale

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

Impostare l'ambiente TFF

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

import nest_asyncio
nest_asyncio.apply()

Definisci il modello da addestrare

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

Configura gli esecutori remoti

Per impostazione predefinita, TFF esegue tutti i calcoli in locale. In questo passaggio diciamo a TFF di connettersi ai servizi Kubernetes che abbiamo impostato sopra. Assicurati di copiare qui l'indirizzo IP del tuo servizio.

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)

Esegui allenamento

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