#Demo :基于联邦学习的图像识别
from __ future__ import absolute_import, division, print_function
import collections
import warnings
from six.moves import range
import numpy as np
import six
import tensorflow as tf
import tensorflow_federated as tff
让我们从数据开始。联合学习需要联合数据集,即来自多个用户的数据集合。联合数据通常是非独立的,这带来了一系列独特的挑战。
为了便于进行实验,我们在TFF存储库中注入了一些数据集,其中包括MNIST的联合版本,它包含已使用Leaf重新处理过的原始NIST数据集的版本,以便保证数据由原始作者输入产生。由于每个作者都有独特的风格,因此该数据集展现了联合数据集所期望的非独立同分布。
这是我们如何加载它。
#load_data()返回的数据集是 tff.simulation.ClientData类的实例,该tff.simulation.ClientData类允许您枚举用户集,构造一个代表特定用户数据的tf.data.Dataset类,查询个体元素结构。
emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data()
以下是使用该tff.simulation.ClientData类浏览数据集内容的方法。请记住,尽管此界面允许您遍历客户端ID,但这只是模拟数据的功能。如您将很快看到的,联合学习框架不使用客户身份-它们的唯一目的是允许您选择数据的子集进行仿真。
print(len(emnist_train.client_ids))
emnist_train.element_type_structure
example_dataset = emnist_train.create_tf_dataset_for_client(emnist_train.client_ids[0]) #返回 tf.data.Dataset 对象
example_element = iter(example_dataset).next()
example_element['label'].numpy()
from matplotlib import pyplot as plt
plt.imshow(example_element['pixels'].numpy(), cmap='gray', aspect='equal')
plt.grid('off')
_ = plt.show()
由于数据已经是一个tf.data.Dataset类,因此可以使用数据集转换来完成预处理。在这里,我们扁平化28x28图像转换成784-元素阵列,对样本进行打乱,把它们组织成批次,并重命名特征pixels,和label 为x和 y(在Keras使用)。我们还会在据集中添加一个repeat函数数以运行多个epoch。
NUM_CLIENTS = 10
NUM_EPOCHS = 10
BATCH_SIZE = 20
SHUFFLE_BUFFER = 500
def preprocess(dataset):
def element_fn(element):
return collections.OrderedDict([
('x', tf.reshape(element['pixels'], [-1])),
('y', tf.reshape(element['label'], [1])),
])
return dataset.repeat(NUM_EPOCHS).map(element_fn).shuffle(
SHUFFLE_BUFFER).batch(BATCH_SIZE)
preprocessed_example_dataset = preprocess(example_dataset)
sample_batch = tf.nest.map_structure(lambda x: x.numpy(), iter(preprocessed_example_dataset).next())
我们几乎已具备构建联邦数据集的所有构建块。在模拟中将联合数据馈送给TFF的方法之一就是简单地作为Python列表,列表中的每个元素都保存单个用户的数据,无论是列表还是tf.data.Dataset。因为我们已经有一个提供后者的接口,所以让我们使用它。这是一个简单的帮助程序功能,该功能将从给定的用户集中构造数据集列表,作为一轮培训或评估的输入。
def make_federated_data(client_data, client_ids):
return [preprocess(client_data.create_tf_dataset_for_client(x))
for x in client_ids]
现在,我们如何选择客户?
在典型的联合培训场景中,我们正在处理大量潜在的用户设备,其中只有一小部分可在给定的时间点进行培训。例如,当客户端设备是仅在插入电源,从计量网络断开或处于空闲状态时才参与培训的移动电话时,就是这种情况。 当然,我们处于仿真环境中,所有数据都在本地可用。通常,然后,在运行模拟时,我们将简单地抽样要参与每一轮培训的客户的随机子集,通常每一轮都不同。 就是说,正如您通过研究有关联合平均算法的论文所发现的那样, 在每个回合中具有随机采样的客户子集的系统中实现收敛可能需要 一段时间,并且在其中进行数百回合是不切实际的。本互动教程。 相反,我们要做的是对一组客户端进行一次采样,并在各回合中重复使用同一组客户端,以加快收敛速度(有意过分适应这几位用户的数据)。我们将其作为练习,供读者修改本教程以模拟随机抽样-这相当容易做到(一旦这样做,请记住,使模型收敛可能要花一些时间)。
sample_clients = emnist_train.client_ids[0:NUM_CLIENTS]
federated_train_data = make_federated_data(emnist_train, sample_clients)
如果您正在使用Keras,则您可能已经具有构造Keras模型的代码。这是一个满足我们需求的简单模型的示例。
def create_compiled_keras_model():
model = tf.keras.models.Sequential([
tf.keras.layers.Dense(
10, activation=tf.nn.softmax, kernel_initializer='zeros', input_shape=(784,))])
model.compile(
loss=tf.keras.losses.SparseCategoricalCrossentropy(),
optimizer=tf.keras.optimizers.SGD(learning_rate=0.02),
metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])
return model
关于compile的一个重要笔记。如下所述,在联合平均算法中使用时,optimizer仅为总优化算法的一半, 因为它仅用于计算每个客户端上的本地模型更新。该算法的其余部分涉及如何在客户端上平均这些更新,以及如何将它们应用于服务器上的全局模型。特别是,这意味着此处使用的优化器和学习率的选择可能需要与用于在标准iid数据集上训练模型的选择不同。我们建议从常规SGD开始,学习速度可能会比平常小。我们在这里使用的学习率尚未经过精心调整,可以随时尝试。 为了将任何模型与TFF一起使用,需要将其包装在tff.learning.Model接口的实例中,与Keras相似,该接口实例公开用于标记模型的前向通过,元数据属性等的方法,但还引入了其他元素,例如方法控制计算联邦指标的过程。我们暂时不用担心。如果您具有上面刚刚定义的已编译Keras模型,则可以通过调用 tff.learning.from_compiled_keras_model,将模型和示例数据批处理作为参数传递,让TFF为您包装它,如下所示。
def model_fn():
keras_model = create_compiled_keras_model()
return tff.learning.from_compiled_keras_model(keras_model, sample_batch)
现在,我们已经包装了一个tff.learning.Model与TFF一起使用的模型,可以通过调用helper函数tff.learning.build_federated_averaging_process,让TFF构造一个Federated Averaging算法,如下所示。
请记住,该参数必须是构造函数(例如model_fn 上述),而不是已经构造的实例,以便可以在TFF控制的上下文中进行模型的构造(如果您对此原因感到好奇),我们建议您阅读有关自定义算法的后续教程 )。
iterative_process = tff.learning.build_federated_averaging_process(model_fn)
刚刚发生了什么?TFF构造了一对联合计算,并将它们打包为一个tff.utils.IterativeProcess,其中这些计算
可以在一对属性initialize和next使用。简而言之,联邦计算是以TFF的内部语言编写的程序,可以表示多种联邦算法(您可以在自定义算法教程中找到有关此的更多信息)。在这种情况下,这两个计算生成并打包到iterative_process中 实现Federated Averaging。 TFF的目标是定义计算并可以在真实的联邦学习设置中执行,但是目前仅实现本地执行仿真运行。要在模拟器中执行计算,您只需像Python函数一样调用它即可。这个默认的解释环境不是为高性能而设计的,但足以满足本教程的要求。我们希望在未来的版本中提供更高性能的仿真运行时,以促进大规模研究。
让我们从第一组联邦计算initialize计算开始。与所有联合计算一样,您可以将其视为一个函数。该计算不带参数,并返回一个结果-服务器上联合平均进程状态的表示。尽管我们不想深入了解TFF的细节,但了解这种状态看起来可能是有益的。您可以将其可视化如下。
str(iterative_process.initialize.type_signature)
而上述类型签名乍一看起来有点神秘,你可以认识到服务器状态包括一个模型(MNIST的初始模型参数,将分发给所有设备),和optimizer_state(附加信息维护的服务器,如用于超参数时间表round的数量等等)。
让我们调用初始化计算来构造服务器状态。
state = iterative_process.initialize()
第二组联邦计算,next,代表了一轮联邦平均,它包括将服务器状态(包括模型参数)推给客户机,对它们的本地数据进行设备上的培训,收集和平均模型更新,并在服务器上生成一个新的更新模型。
从概念上讲,可以认为next具有如下所示的函数类型签名。
SERVER_STATE, FEDERATED_DATA -> SERVER_STATE, TRAINING_METRICS
特别是,一个人应该考虑next()不作为一个运行在一个服务器上函数,而是作为整个分散式计算的一个声明性函数表示——服务器提供一些输入(SERVER_STATE),但每个参与设备贡献自己的本地数据集。
让我们进行一轮训练并将结果可视化。我们可以使用前面为用户示例生成的联邦数据。
state, metrics = iterative_process.next(state, federated_train_data)
print('round 1, metrics={}'.format(metrics))
让我们再跑几round。正如前面提到的,通常在这一点上为了模拟实际部署用户不断地来来去去的场景,每一轮你会从一个新的随机被选的用户样本中选择一个模拟数据的子集;但在这个互动的笔记本,为了演示我们就重用相同的用户,保证系统快速收敛。
for round_num in range(2, 11):
state, metrics = iterative_process.next(state, federated_train_data)
print('round {:2d}, metrics={}'.format(round_num, metrics))
每轮联合训练后训练损失减小,表明模型收敛。但是,这些培训指标有一些重要的注意事项,请参阅本教程后面的评估部分。
Keras是TensorFlow推荐的高级模型API,我们鼓励尽可能在TFF中使用Keras模型(通过TFF.learning.from_keras_model或TFF.learning.from_compiled_keras_model)。然而,tff.learning提供了一个较低级的模型接口,tff.learning.Model。这暴露了使用模型用于联邦学习所需的最小功能。直接实现这个接口(可能仍然使用类似tf.keras.layers这样的构建块)允许最大程度的自定义,而无需修改联邦学习算法的内部机制。让我们从头再做一遍。
第一步是识别我们要使用的TensorFlow变量。为了使代码更加清晰,让我们定义一个数据结构来表示整个集。这包括我们将训练变量如权重和阈值,以及各种训练期间将更新的累积统计信息和计数器变量,如loss_sum
accuracy_sum, num_examples。
MnistVariables = collections.namedtuple('MnistVariables', 'weights bias num_examples loss_sum accuracy_sum')
下面是一个创建变量的方法。为了简单起见,我们将所有统计数据表示为tf.float32。因为这将在稍后阶段消除类型转换的需要。将变量初始化器包装成lambdas是资源变量强加的一个要求。
def create_mnist_variables():
return MnistVariables(
weights = tf.Variable(
lambda: tf.zeros(dtype=tf.float32, shape=(784, 10)),
name='weights',
trainable=True),
bias = tf.Variable(
lambda: tf.zeros(dtype=tf.float32, shape=(10)),
name='bias',
trainable=True),
num_examples = tf.Variable(0.0, name='num_examples', trainable=False),
loss_sum = tf.Variable(0.0, name='loss_sum', trainable=False),
accuracy_sum = tf.Variable(0.0, name='accuracy_sum', trainable=False))
有了模型参数和累积统计数据的变量,我们现在可以定义前向传递方法,该方法计算损失、发出预测并更新单批输入数据的累积统计数据,如下所示。
def mnist_forward_pass(variables, batch):
y = tf.nn.softmax(tf.matmul(batch['x'], variables.weights) + variables.bias)
predictions = tf.cast(tf.argmax(y, 1), tf.int32)
flat_labels = tf.reshape(batch['y'], [-1])
loss = -tf.reduce_mean(tf.reduce_sum(
tf.one_hot(flat_labels, 10) * tf.math.log(y), axis=[1]))
accuracy = tf.reduce_mean(
tf.cast(tf.equal(predictions, flat_labels), tf.float32))
num_examples = tf.cast(tf.size(batch['y']), tf.float32)
variables.num_examples.assign_add(num_examples)
variables.loss_sum.assign_add(loss * num_examples)
variables.accuracy_sum.assign_add(accuracy * num_examples)
return loss, predictions
接下来,我们定义一个返回一组本地度量的函数,同样使用TensorFlow。这些值(除了自动处理的模型更新之外)有资格在联邦学习或评估过程中聚合到服务器。
在这里,我们简单地返回平均损失和准确性,以及
num_examples,在计算联合聚合时,我们需要num_examples正确地权衡来自不同用户的贡献。
def get_local_mnist_metrics(variables):
return collections.OrderedDict([
('num_examples', variables.num_examples),
('loss', variables.loss_sum / variables.num_examples),
('accuracy', variables.accuracy_sum / variables.num_examples)
])
最后,我们需要确定如何通过get_local_mnist_metrics聚合每个设备发出的本地度量。这是代码中唯一没有使用TensorFlow编写的部分——它是用TFF表示的联合计算。如果你想深入研究,浏览一下自定义算法教程,但在大多数应用程序中,你不需要这样做;下面所示模式的变体应该足够了。它是这样的:
@tff.federated_computation
def aggregate_mnist_metrics_across_clients(metrics):
return {
'num_examples': tff.federated_sum(metrics.num_examples),
'loss': tff.federated_mean(metrics.loss, metrics.num_examples),
'accuracy': tff.federated_mean(metrics.accuracy, metrics.num_examples)
}
input metrics参数对应于上面get_local_mnist_metrics返回的OrderedDict,但关键的是这些值不再是tf。张量——它们作为tff被“框起来”。值,说明您不能再使用TensorFlow操作它们,而只能使用TFF的联邦操作符(如TFF)。federated_mean tff.federated_sum。返回的全局聚合字典定义了在服务 器上可用的一组指标。
有了以上所有内容,我们就可以构建一个用于TFF的模型表示,类似于让TFF摄取Keras模型时为您生成的模型。
class MnistModel(tff.learning.Model):
def __init__(self):
self._variables = create_mnist_variables()
@property
def trainable_variables(self):
return [self._variables.weights, self._variables.bias]
@property
def non_trainable_variables(self):
return []
@property
def local_variables(self):
return [
self._variables.num_examples, self._variables.loss_sum,
self._variables.accuracy_sum
]
@property
def input_spec(self):
return collections.OrderedDict([('x', tf.TensorSpec([None, 784],
tf.float32)),
('y', tf.TensorSpec([None, 1], tf.int32))])
@tf.function
def forward_pass(self, batch, training=True):
del training
loss, predictions = mnist_forward_pass(self._variables, batch)
num_exmaples = tf.shape(batch['x'])[0]
return tff.learning.BatchOutput(
loss=loss, predictions=predictions, num_examples=num_exmaples)
@tf.function
def report_local_outputs(self):
return get_local_mnist_metrics(self._variables)
@property
def federated_output_computation(self):
return aggregate_mnist_metrics_across_clients