การจำลองประสิทธิภาพสูงด้วย TFF

บทช่วยสอนนี้จะอธิบายวิธีตั้งค่าการจำลองประสิทธิภาพสูงด้วย TFF ในสถานการณ์ทั่วไปที่หลากหลาย

สิ่งที่ต้องทำ(b/134543154): เติมเนื้อหา บางสิ่งที่จะกล่าวถึงที่นี่:

• โดยใช้ GPU ในการตั้งค่าเครื่องเดียว
• การตั้งค่าเครื่องหลายเครื่องบน GCP/GKE โดยมีและไม่มี TPU
• เชื่อมต่อแบ็กเอนด์เหมือน MapReduce
• ข้อจำกัดในปัจจุบันและเมื่อใด/อย่างไรจะผ่อนคลาย

ก่อนที่เราจะเริ่มต้น

ขั้นแรก ตรวจสอบให้แน่ใจว่าโน้ตบุ๊กของคุณเชื่อมต่อกับแบ็กเอนด์ที่มีคอมไพล์ส่วนประกอบที่เกี่ยวข้อง (รวมถึงการขึ้นต่อกันของ gRPC สำหรับสถานการณ์แบบหลายเครื่อง)

มาเริ่มด้วยการโหลดตัวอย่าง MNIST จากเว็บไซต์ TFF และประกาศฟังก์ชัน Python ที่จะเรียกใช้การทดลองวนรอบเล็กๆ ในกลุ่มลูกค้า 10 ราย

!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).shuffle(500).batch(20).map(map_fn)


train_data = [client_data(n) for n in range(10)]
element_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=element_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 _ in range(num_rounds):
    t1 = time.time()
    state, metrics = trainer.next(state, train_data)
    t2 = time.time()
    print('metrics {m}, round time {t:.2f} seconds'.format(
        m=metrics, t=t2 - t1))

การจำลองด้วยเครื่องเดียว

ตอนนี้โดยค่าเริ่มต้น

evaluate()

metrics <sparse_categorical_accuracy=0.13858024775981903,loss=3.0073554515838623>, round time 3.59 seconds
metrics <sparse_categorical_accuracy=0.1796296238899231,loss=2.749046802520752>, round time 2.29 seconds
metrics <sparse_categorical_accuracy=0.21656379103660583,loss=2.514779567718506>, round time 2.33 seconds
metrics <sparse_categorical_accuracy=0.2637860178947449,loss=2.312587261199951>, round time 2.06 seconds
metrics <sparse_categorical_accuracy=0.3334362208843231,loss=2.068122386932373>, round time 2.00 seconds
metrics <sparse_categorical_accuracy=0.3737654387950897,loss=1.9268712997436523>, round time 2.42 seconds
metrics <sparse_categorical_accuracy=0.4296296238899231,loss=1.7216310501098633>, round time 2.20 seconds
metrics <sparse_categorical_accuracy=0.4655349850654602,loss=1.6489890813827515>, round time 2.18 seconds
metrics <sparse_categorical_accuracy=0.5048353672027588,loss=1.5485210418701172>, round time 2.16 seconds
metrics <sparse_categorical_accuracy=0.5564814805984497,loss=1.4140453338623047>, round time 2.41 seconds

การจำลองหลายเครื่องบน GCP/GKE, GPU, TPU และอื่นๆ...

มาเร็ว ๆ นี้