Simulação de alto desempenho com Kubernetes

Este tutorial descreverá como configurar a simulação de alto desempenho usando um tempo de execução TFF em execução no Kubernetes. O modelo é o mesmo que no tutorial anterior, simulações de alto desempenho com TFF. A única diferença é que aqui usamos um pool de trabalhadores em vez de um executor local.

Este tutorial refere-se a do Google Cloud GKE para criar o cluster Kubernetes, mas todos os passos depois que o cluster é criado pode ser usado com qualquer instalação Kubernetes.

Ver no TensorFlow.org Executar no Google Colab Ver fonte no GitHub Baixar caderno

Inicie os TFF Workers no GKE

Crie um cluster do Kubernetes

A etapa a seguir só precisa ser executada uma vez. O cluster pode ser reutilizado para cargas de trabalho futuras.

Siga as instruções GKE para criar um cluster de contêiner . O restante deste tutorial assume que o cluster é chamado tff-cluster , mas o nome real não é importante. Parar de seguir as instruções quando você chegar ao "Passo 5: Implantar sua aplicação".

Implantar o aplicativo de trabalho TFF

Os comandos para interagir com GCP podem ser executados localmente ou na nuvem Shell Google . Recomendamos o Google Cloud Shell, pois não requer configuração adicional.

  1. Execute o seguinte comando para iniciar o aplicativo Kubernetes.
$ kubectl create deployment tff-workers --image=gcr.io/tensorflow-federated/remote-executor-service:latest
  1. Adicione um balanceador de carga para o aplicativo.
$ kubectl expose deployment tff-workers --type=LoadBalancer --port 80 --target-port 8000

Procure o endereço IP do balanceador de carga no Google Cloud Console. Você precisará dele mais tarde para conectar o loop de treinamento ao aplicativo de trabalho.

(Como alternativa) Inicie o contêiner do Docker localmente

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

Configurar ambiente TFF

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

import nest_asyncio
nest_asyncio.apply()

Defina o modelo para treinar

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 os executores remotos

Por padrão, o TFF executa todos os cálculos localmente. Nesta etapa, instruímos a TFF a se conectar aos serviços do Kubernetes configurados acima. Certifique-se de copiar o endereço IP do seu serviço aqui.

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)

Executar o treinamento

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