Mô phỏng hiệu suất cao với Kubernetes

Hướng dẫn này sẽ mô tả cách thiết lập mô phỏng hiệu suất cao bằng thời gian chạy TFF chạy trên Kubernetes. Mô hình này cũng giống như trong hướng dẫn trước, mô phỏng cao hiệu suất với TFF. Sự khác biệt duy nhất là ở đây chúng tôi sử dụng một nhóm công nhân thay vì một trình thực thi cục bộ.

Hướng dẫn này đề cập đến Google Cloud của GKE để tạo ra các cụm Kubernetes, nhưng tất cả các bước sau cụm được tạo ra có thể được sử dụng với bất kỳ cài đặt Kubernetes.

Xem trên TensorFlow.org Chạy trong Google Colab Xem nguồn trên GitHub Tải xuống sổ ghi chép

Khởi chạy Công nhân TFF trên GKE

Tạo một cụm Kubernetes

Bước sau chỉ cần thực hiện một lần. Cụm có thể được sử dụng lại cho các khối lượng công việc trong tương lai.

Làm theo hướng dẫn GKE để tạo một cluster chứa . Phần còn lại của hướng dẫn này giả định rằng các cụm được đặt tên tff-cluster , nhưng tên thực tế là không quan trọng. Dừng theo các hướng dẫn khi bạn nhận được để "Bước 5: Triển khai ứng dụng của bạn".

Triển khai Ứng dụng Nhân viên TFF

Các lệnh để tương tác với GCP có thể chạy cục bộ hoặc trong Google Cloud Shell . Chúng tôi khuyên bạn nên sử dụng Google Cloud Shell vì nó không yêu cầu thiết lập thêm.

  1. Chạy lệnh sau để khởi chạy ứng dụng Kubernetes.
$ kubectl create deployment tff-workers --image=gcr.io/tensorflow-federated/remote-executor-service:latest
  1. Thêm bộ cân bằng tải cho ứng dụng.
$ kubectl expose deployment tff-workers --type=LoadBalancer --port 80 --target-port 8000

Tra cứu địa chỉ IP của bộ cân bằng tải trên Google Cloud Console. Sau này, bạn sẽ cần nó để kết nối vòng đào tạo với ứng dụng worker.

(Cách khác) Khởi chạy Docker Container cục bộ

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

Thiết lập môi trường TFF

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

import nest_asyncio
nest_asyncio.apply()

Xác định mô hình để đào tạo

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

Thiết lập người thực thi từ xa

Theo mặc định, TFF thực thi tất cả các tính toán cục bộ. Trong bước này, chúng tôi yêu cầu TFF kết nối với các dịch vụ Kubernetes mà chúng tôi đã thiết lập ở trên. Đảm bảo sao chép địa chỉ IP của dịch vụ của bạn tại đây.

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)

Chạy đào tạo

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