Simulación de alto rendimiento con Kubernetes

Este tutorial describirá cómo configurar una simulación de alto rendimiento mediante un tiempo de ejecución de TFF que se ejecuta en Kubernetes. El modelo es el mismo que en el tutorial anterior, las simulaciones de alto rendimiento con TFF. La única diferencia es que aquí usamos un grupo de trabajadores en lugar de un ejecutor local.

Este tutorial se refiere a la nube de Google GKE para crear el clúster Kubernetes, pero todos los pasos después del cluster se crea se puede utilizar con cualquier instalación Kubernetes.

Ver en TensorFlow.org Ejecutar en Google Colab Ver fuente en GitHub Descargar cuaderno

Lanzar los trabajadores de TFF en GKE

Crear un clúster de Kubernetes

El siguiente paso solo debe realizarse una vez. El clúster se puede reutilizar para cargas de trabajo futuras.

Siga las instrucciones gke a crear una agrupación de contenedor . El resto de este tutorial asume que el clúster se denomina tff-cluster , pero el nombre real no es importante. Dejar de seguir las instrucciones cuando llegue al "Paso 5: Implementar su aplicación".

Implementar la aplicación TFF Worker

Los comandos para interactuar con la BPC se pueden ejecutar de forma local o en el Shell de Cloud de Google . Recomendamos Google Cloud Shell ya que no requiere configuración adicional.

  1. Ejecute el siguiente comando para iniciar la aplicación Kubernetes.
$ kubectl create deployment tff-workers --image=gcr.io/tensorflow-federated/remote-executor-service:latest
  1. Agregue un equilibrador de carga para la aplicación.
$ kubectl expose deployment tff-workers --type=LoadBalancer --port 80 --target-port 8000

Busque la dirección IP del balanceador de carga en Google Cloud Console. Lo necesitará más tarde para conectar el ciclo de entrenamiento a la aplicación del trabajador.

(Alternativamente) Inicie el contenedor Docker localmente

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

Configurar el entorno TFF

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

import nest_asyncio
nest_asyncio.apply()

Definir el modelo a entrenar

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

Configurar los ejecutores remotos

De forma predeterminada, TFF ejecuta todos los cálculos localmente. En este paso, le decimos a TFF que se conecte a los servicios de Kubernetes que configuramos anteriormente. Asegúrese de copiar la dirección IP de su servicio aquí.

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)

Ejecutar entrenamiento

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