 TensorFlow.org에서 보기 |  Google Colab에서 실행 |  GitHub에서 소스 보기 |  노트북 다운로드 |
이것은 대안입니다 빌드 자신의 연합 학습 알고리즘 튜토리얼과 simple_fedavg 에 대한 사용자 정의 반복적 인 프로세스 구축 예 연합 평균 알고리즘을. 이 튜토리얼은 사용 TFF 옵티 마이저 대신 Keras 옵티 마이저를. TFF 옵티마이저 추상화는 TFF 반복 프로세스에 통합하기 쉽도록 state-in-state-out으로 설계되었습니다. tff.learning API는 또한 입력 인수로 TFF 옵티 마이저에 동의합니다.
시작하기 전에
시작하기 전에 다음을 실행하여 환경이 올바르게 설정되었는지 확인하십시오. 당신이 인사말을 표시되지 않는 경우를 참조하십시오 설치 지침은 가이드.
!pip install --quiet --upgrade tensorflow-federated-nightly
!pip install --quiet --upgrade nest-asyncio
import nest_asyncio
nest_asyncio.apply()
import functools
import attr
import numpy as np
import tensorflow as tf
import tensorflow_federated as tff
데이터 및 모델 준비
EMNIST 데이터 처리 모델은 매우 유사 simple_fedavg 예.
only_digits=True
# Load dataset.
emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data(only_digits)
# Define preprocessing functions.
def preprocess_fn(dataset, batch_size=16):
def batch_format_fn(element):
return (tf.expand_dims(element['pixels'], -1), element['label'])
return dataset.batch(batch_size).map(batch_format_fn)
# Preprocess and sample clients for prototyping.
train_client_ids = sorted(emnist_train.client_ids)
train_data = emnist_train.preprocess(preprocess_fn)
central_test_data = preprocess_fn(
emnist_train.create_tf_dataset_for_client(train_client_ids[0]))
# Define model.
def create_keras_model():
"""The CNN model used in https://arxiv.org/abs/1602.05629."""
data_format = 'channels_last'
input_shape = [28, 28, 1]
max_pool = functools.partial(
tf.keras.layers.MaxPooling2D,
pool_size=(2, 2),
padding='same',
data_format=data_format)
conv2d = functools.partial(
tf.keras.layers.Conv2D,
kernel_size=5,
padding='same',
data_format=data_format,
activation=tf.nn.relu)
model = tf.keras.models.Sequential([
conv2d(filters=32, input_shape=input_shape),
max_pool(),
conv2d(filters=64),
max_pool(),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(512, activation=tf.nn.relu),
tf.keras.layers.Dense(10 if only_digits else 62),
])
return model
# Wrap as `tff.learning.Model`.
def model_fn():
keras_model = create_keras_model()
return tff.learning.from_keras_model(
keras_model,
input_spec=central_test_data.element_spec,
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True))
맞춤형 반복 프로세스
많은 경우 연합 알고리즘에는 4가지 주요 구성 요소가 있습니다.
- 서버-클라이언트 브로드캐스트 단계.
- 로컬 클라이언트 업데이트 단계.
- 클라이언트에서 서버로의 업로드 단계.
- 서버 업데이트 단계입니다.
TFF, 우리는 일반적으로 같은 연합 알고리즘 나타내는 tff.templates.IterativeProcess (우리는 단지로 참조 IterativeProcess 걸쳐 참조). 이것은 포함하는 클래스 initialize 하고 next 기능을. 여기서, initialize 서버를 초기화하는데 사용되며, next 연합 알고리즘의 하나의 통신 라운드를 수행한다.
클라이언트 업데이트 단계에서 옵티마이저를 사용하고 서버 업데이트 단계에서 또 다른 옵티마이저를 사용하는 연합 평균(FedAvg) 알고리즘을 구축하기 위해 다양한 구성 요소를 소개합니다. 클라이언트 및 서버 업데이트의 핵심 로직은 순수 TF 블록으로 표현될 수 있습니다.
TF 블록: 클라이언트 및 서버 업데이트
각 클라이언트에서 로컬 client_optimizer 초기화 및 클라이언트 모델의 가중치를 업데이트하는 데 사용됩니다. 서버에서 server_optimizer 이전 라운드에서 상태를 사용하고, 다음 라운드에 대한 상태를 업데이트합니다.
@tf.function
def client_update(model, dataset, server_weights, client_optimizer):
"""Performs local training on the client's dataset."""
# Initialize the client model with the current server weights.
client_weights = model.trainable_variables
# Assign the server weights to the client model.
tf.nest.map_structure(lambda x, y: x.assign(y),
client_weights, server_weights)
# Initialize the client optimizer.
trainable_tensor_specs = tf.nest.map_structure(
lambda v: tf.TensorSpec(v.shape, v.dtype), client_weights)
optimizer_state = client_optimizer.initialize(trainable_tensor_specs)
# Use the client_optimizer to update the local model.
for batch in iter(dataset):
with tf.GradientTape() as tape:
# Compute a forward pass on the batch of data.
outputs = model.forward_pass(batch)
# Compute the corresponding gradient.
grads = tape.gradient(outputs.loss, client_weights)
# Apply the gradient using a client optimizer.
optimizer_state, updated_weights = client_optimizer.next(
optimizer_state, client_weights, grads)
tf.nest.map_structure(lambda a, b: a.assign(b),
client_weights, updated_weights)
# Return model deltas.
return tf.nest.map_structure(tf.subtract, client_weights, server_weights)
@attr.s(eq=False, frozen=True, slots=True)
class ServerState(object):
trainable_weights = attr.ib()
optimizer_state = attr.ib()
@tf.function
def server_update(server_state, mean_model_delta, server_optimizer):
"""Updates the server model weights."""
# Use aggregated negative model delta as pseudo gradient.
negative_weights_delta = tf.nest.map_structure(
lambda w: -1.0 * w, mean_model_delta)
new_optimizer_state, updated_weights = server_optimizer.next(
server_state.optimizer_state, server_state.trainable_weights,
negative_weights_delta)
return tff.structure.update_struct(
server_state,
trainable_weights=updated_weights,
optimizer_state=new_optimizer_state)
TFF 블록 : tff.tf_computation 및 tff.federated_computation
이제 오케스트레이션에 TFF를 사용하고 FedAvg에 대한 반복 프로세스를 구축합니다. 우리는 위에서 정의 된 TF 블록 포장해야 tff.tf_computation 및 사용 TFF 방법 tff.federated_broadcast , tff.federated_map , tff.federated_mean A의 tff.federated_computation 기능. 사용하기 쉬운 tff.learning.optimizers.Optimizer 와 API를 initialize 하고 next 사용자 정의 반복적 인 프로세스를 정의 할 때 기능.
# 1. Server and client optimizer to be used.
server_optimizer = tff.learning.optimizers.build_sgdm(
learning_rate=0.05, momentum=0.9)
client_optimizer = tff.learning.optimizers.build_sgdm(
learning_rate=0.01)
# 2. Functions return initial state on server.
@tff.tf_computation
def server_init():
model = model_fn()
trainable_tensor_specs = tf.nest.map_structure(
lambda v: tf.TensorSpec(v.shape, v.dtype), model.trainable_variables)
optimizer_state = server_optimizer.initialize(trainable_tensor_specs)
return ServerState(
trainable_weights=model.trainable_variables,
optimizer_state=optimizer_state)
@tff.federated_computation
def server_init_tff():
return tff.federated_value(server_init(), tff.SERVER)
# 3. One round of computation and communication.
server_state_type = server_init.type_signature.result
print('server_state_type:\n',
server_state_type.formatted_representation())
trainable_weights_type = server_state_type.trainable_weights
print('trainable_weights_type:\n',
trainable_weights_type.formatted_representation())
# 3-1. Wrap server and client TF blocks with `tff.tf_computation`.
@tff.tf_computation(server_state_type, trainable_weights_type)
def server_update_fn(server_state, model_delta):
return server_update(server_state, model_delta, server_optimizer)
whimsy_model = model_fn()
tf_dataset_type = tff.SequenceType(whimsy_model.input_spec)
print('tf_dataset_type:\n',
tf_dataset_type.formatted_representation())
@tff.tf_computation(tf_dataset_type, trainable_weights_type)
def client_update_fn(dataset, server_weights):
model = model_fn()
return client_update(model, dataset, server_weights, client_optimizer)
# 3-2. Orchestration with `tff.federated_computation`.
federated_server_type = tff.FederatedType(server_state_type, tff.SERVER)
federated_dataset_type = tff.FederatedType(tf_dataset_type, tff.CLIENTS)
@tff.federated_computation(federated_server_type, federated_dataset_type)
def run_one_round(server_state, federated_dataset):
# Server-to-client broadcast.
server_weights_at_client = tff.federated_broadcast(
server_state.trainable_weights)
# Local client update.
model_deltas = tff.federated_map(
client_update_fn, (federated_dataset, server_weights_at_client))
# Client-to-server upload and aggregation.
mean_model_delta = tff.federated_mean(model_deltas)
# Server update.
server_state = tff.federated_map(
server_update_fn, (server_state, mean_model_delta))
return server_state
# 4. Build the iterative process for FedAvg.
fedavg_process = tff.templates.IterativeProcess(
initialize_fn=server_init_tff, next_fn=run_one_round)
print('type signature of `initialize`:\n',
fedavg_process.initialize.type_signature.formatted_representation())
print('type signature of `next`:\n',
fedavg_process.next.type_signature.formatted_representation())
server_state_type:
<
trainable_weights=<
float32[5,5,1,32],
float32[32],
float32[5,5,32,64],
float32[64],
float32[3136,512],
float32[512],
float32[512,10],
float32[10]
>,
optimizer_state=<
float32[5,5,1,32],
float32[32],
float32[5,5,32,64],
float32[64],
float32[3136,512],
float32[512],
float32[512,10],
float32[10]
>
>
trainable_weights_type:
<
float32[5,5,1,32],
float32[32],
float32[5,5,32,64],
float32[64],
float32[3136,512],
float32[512],
float32[512,10],
float32[10]
>
tf_dataset_type:
<
float32[?,28,28,1],
int32[?]
>*
type signature of `initialize`:
( -> <
trainable_weights=<
float32[5,5,1,32],
float32[32],
float32[5,5,32,64],
float32[64],
float32[3136,512],
float32[512],
float32[512,10],
float32[10]
>,
optimizer_state=<
float32[5,5,1,32],
float32[32],
float32[5,5,32,64],
float32[64],
float32[3136,512],
float32[512],
float32[512,10],
float32[10]
>
>@SERVER)
type signature of `next`:
(<
server_state=<
trainable_weights=<
float32[5,5,1,32],
float32[32],
float32[5,5,32,64],
float32[64],
float32[3136,512],
float32[512],
float32[512,10],
float32[10]
>,
optimizer_state=<
float32[5,5,1,32],
float32[32],
float32[5,5,32,64],
float32[64],
float32[3136,512],
float32[512],
float32[512,10],
float32[10]
>
>@SERVER,
federated_dataset={<
float32[?,28,28,1],
int32[?]
>*}@CLIENTS
> -> <
trainable_weights=<
float32[5,5,1,32],
float32[32],
float32[5,5,32,64],
float32[64],
float32[3136,512],
float32[512],
float32[512,10],
float32[10]
>,
optimizer_state=<
float32[5,5,1,32],
float32[32],
float32[5,5,32,64],
float32[64],
float32[3136,512],
float32[512],
float32[512,10],
float32[10]
>
>@SERVER)
알고리즘 평가
중앙 집중식 평가 데이터 세트에서 성능을 평가합니다.
def evaluate(server_state):
keras_model = create_keras_model()
tf.nest.map_structure(
lambda var, t: var.assign(t),
keras_model.trainable_weights, server_state.trainable_weights)
metric = tf.keras.metrics.SparseCategoricalAccuracy()
for batch in iter(central_test_data):
preds = keras_model(batch[0], training=False)
metric.update_state(y_true=batch[1], y_pred=preds)
return metric.result().numpy()
server_state = fedavg_process.initialize()
acc = evaluate(server_state)
print('Initial test accuracy', acc)
# Evaluate after a few rounds
CLIENTS_PER_ROUND=2
sampled_clients = train_client_ids[:CLIENTS_PER_ROUND]
sampled_train_data = [
train_data.create_tf_dataset_for_client(client)
for client in sampled_clients]
for round in range(20):
server_state = fedavg_process.next(server_state, sampled_train_data)
acc = evaluate(server_state)
print('Test accuracy', acc)
Initial test accuracy 0.09677419 Test accuracy 0.13978495