Высокопроизводительное моделирование с Kubernetes

В этом руководстве описывается, как настроить высокопроизводительное моделирование с использованием среды выполнения TFF, работающей в Kubernetes. Модель такой же , как и в предыдущем учебнике, моделирование высокой производительности при использовании TFF. Единственная разница в том, что здесь мы используем рабочий пул вместо локального исполнителя.

Это руководство относится к Google Облако GKE создать кластер Kubernetes, но все шаги после кластера создается может быть использован с любой установкой Kubernetes.

Посмотреть на TensorFlow.org Запускаем в Google Colab Посмотреть исходный код на GitHub Скачать блокнот

Запустите TFF Workers на GKE

Создайте кластер Kubernetes

Следующий шаг нужно сделать только один раз. Кластер можно повторно использовать для будущих рабочих нагрузок.

Следуйте инструкции GKE для создания контейнера кластера . Остальная часть этого урока предполагает , что кластер называется tff-cluster , но фактическое имя не имеет значения. Стоп , следуя инструкции , когда вы получаете «Шаг 5: Развертывание приложения».

Развертывание рабочего приложения TFF

Команды для взаимодействия с GCP может работать локально или в Google Cloud Shell . Мы рекомендуем Google Cloud Shell, так как он не требует дополнительной настройки.

  1. Выполните следующую команду, чтобы запустить приложение Kubernetes.
$ kubectl create deployment tff-workers --image=gcr.io/tensorflow-federated/remote-executor-service:latest
  1. Добавьте балансировщик нагрузки для приложения.
$ kubectl expose deployment tff-workers --type=LoadBalancer --port 80 --target-port 8000

Найдите IP-адрес балансировщика нагрузки в Google Cloud Console. Он понадобится вам позже, чтобы подключить цикл обучения к рабочему приложению.

(Или же) Запустите контейнер Docker локально

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

Настроить среду TFF

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

import nest_asyncio
nest_asyncio.apply()

Определите модель для обучения

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

Настройка удаленных исполнителей

По умолчанию TFF выполняет все вычисления локально. На этом этапе мы говорим TFF ​​подключиться к сервисам Kubernetes, которые мы настроили выше. Не забудьте скопировать сюда IP-адрес вашего сервиса.

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)

Выполнить обучение

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