High-performance simulations with TFF

This tutorial will describe how to setup high-performance simulations with TFF in a variety of common scenarios.

TODO(b/134543154): Populate the content, some of the things to cover here:

• using GPUs in a single-machine setup,
• multi-machine setup on GCP/GKE, with and without TPUs,
• interfacing MapReduce-like backends,
• current limitations and when/how they will be relaxed.

View on TensorFlow.org

Run in Google Colab

View source on GitHub

Download notebook

准备工作

首先，请确保您的笔记本电脑连接到具有已编译相关组件（包括用于多机场景的 gRPC 依赖关系）的后端。

现在，我们从加载 TFF 网站上的 MNIST 示例开始，然后声明 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 等上进行多机模拟…

敬请期待。